From 8003ae657ac38496d609404d70cc182a6a25d979 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Tue, 12 May 2026 00:01:37 +0800 Subject: [PATCH] docs: implementation plan for ai_mouse library refactor Expands the 2026-05-11 design spec into ~40 bite-sized tasks across 6 phases (pre-flight golden capture, tools/ extraction, src layout switch, ONNX export, NumPy/ORT rewrite, docs cleanup). Each task is self-contained with full code blocks, exact file paths, and verification commands. TDD where applicable; pure-move tasks use shorter scaffolding. Co-Authored-By: Claude Opus 4.7 --- .../2026-05-11-ai-mouse-library-refactor.md | 4002 +++++++++++++++++ 1 file changed, 4002 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md diff --git a/docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md b/docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md new file mode 100644 index 0000000..7a5c4c7 --- /dev/null +++ b/docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md @@ -0,0 +1,4002 @@ +# ai_mouse Library Refactor Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Refactor `ai_mouse` from a mixed application/library into a slim ONNX-Runtime-based inference SDK (`src/ai_mouse/`), with all training/server/eval code moved to repo-internal `tools/` and pre-trained weights bundled inside the wheel. + +**Architecture:** `src/`-layout package distributable via git URL; runtime depends only on `numpy + onnxruntime`. `tools/` directory holds development-only code (torch, fastapi, etc.). Public API exposes `MouseModel`/`ScrollModel` classes plus cached `generate`/`generate_scroll` helper functions. + +**Tech Stack:** Python 3.12+, NumPy, ONNX Runtime, hatchling (build backend), pytest, uv. Tools-only: PyTorch, FastAPI, scipy, matplotlib. + +**Spec:** [`docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md`](../specs/2026-05-11-ai-mouse-library-design.md) + +--- + +## File Structure (target end state) + +``` +ai_mouse/ (repo root) +├── src/ +│ └── ai_mouse/ # wheel content +│ ├── __init__.py +│ ├── mouse.py +│ ├── scroll.py +│ ├── _coord.py +│ ├── _postprocess.py +│ ├── _assets.py +│ ├── errors.py +│ ├── py.typed +│ └── assets/ +│ ├── flow_model.onnx +│ ├── scroll_decoder.onnx +│ ├── click_dist.json +│ ├── duration_dist.json +│ ├── train_config.json +│ └── scroll_config.json +├── tools/ # dev-only, not in wheel +│ ├── __init__.py +│ ├── __main__.py +│ ├── train.py / serve.py / export_onnx.py +│ ├── trainer.py / models.py / collector.py / config.py +│ ├── server/ / eval/ / data_adapters/ +│ └── scroll/{trainer,models,collector}.py +├── tests/{unit,tools}/ +├── examples/quickstart.py +├── data/ / static/ / docs/ # unchanged +├── pyproject.toml / CHANGELOG.md / README.md / CLAUDE.md +``` + +--- + +## Phase Map + +| Phase | Goal | Validation | +|---|---|---| +| 0 | Capture golden tests + train scroll model | golden npz files committed | +| 1 | Move dev-only code from `ai_mouse/` to `tools/` | `python -m tools train` works; old `from ai_mouse import generate` still works | +| 2 | Switch to `src/` layout + tighten pyproject | `uv build` produces clean wheel; runtime install has no torch | +| 3 | Write ONNX exporter + commit assets | `tools/export_onnx.py` produces `.onnx` files; parity check passes | +| 4 | Rewrite library in NumPy + ORT | Golden tests pass; `import ai_mouse` works without torch | +| 5 | Docs + cleanup | README, CHANGELOG, CLAUDE.md updated; examples runnable | + +--- + +## Phase 0: Pre-flight + +### Task 0.1: Train scroll model + +The repo has `data/scroll_traces.jsonl` but no trained scroll model. The current trainer in `ai_mouse/scroll/trainer.py` exists and works. + +**Files:** +- Read: `ai_mouse/scroll/trainer.py` +- Output: `data/scroll_models/{scroll_model.pt, scroll_config.json}` + +- [ ] **Step 1: Locate the scroll training entry point** + +Run: `uv run python -c "from ai_mouse.scroll.trainer import train; help(train)"` +Confirm there's a callable `train(data_path, output_dir, ...)` with default epochs around 100. + +- [ ] **Step 2: Train the scroll model** + +```bash +uv run python -c " +from pathlib import Path +from ai_mouse.scroll.trainer import train +train( + data_path=Path('data/scroll_traces.jsonl'), + output_dir=Path('data/scroll_models'), +) +" +``` + +Expected: runs ~100 epochs over ~3 minutes on CPU. Loss decreasing. Writes `scroll_model.pt` and `scroll_config.json` to `data/scroll_models/`. + +- [ ] **Step 3: Verify outputs exist** + +Run: `ls data/scroll_models/` +Expected: `scroll_model.pt`, `scroll_config.json`. + +- [ ] **Step 4: Smoke-test inference** + +```bash +uv run python -c " +from ai_mouse.scroll.generator import generate_scroll +events = generate_scroll(0, 1500, mode='target', model_dir='data/scroll_models') +print(f'Got {len(events)} events; sum deltaY = {sum(e[\"deltaY\"] for e in events)}') +" +``` + +Expected: prints something like `Got 14 events; sum deltaY = 1480` (close to 1500). + +- [ ] **Step 5: Commit the model** + +```bash +git add data/scroll_models/scroll_model.pt data/scroll_models/scroll_config.json +git commit -m "chore(scroll): train initial scroll model from scroll_traces.jsonl" +``` + +--- + +### Task 0.2: Build mouse golden npz + +Capture deterministic output from the current torch-based `generate()` for use in regression tests later. + +**Files:** +- Create: `scripts/build_golden_mouse.py` (temporary, will be deleted after Phase 4) +- Output: `tests/unit/data/golden_mouse.npz` + +- [ ] **Step 1: Ensure tests/unit/data/ exists** + +```bash +mkdir -p tests/unit/data +``` + +- [ ] **Step 2: Create the build script** + +Create `scripts/build_golden_mouse.py`: + +```python +"""One-shot script to capture golden mouse trajectories from the current torch +implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite +in Phase 4 produces equivalent output. + +Output: tests/unit/data/golden_mouse.npz +""" +from __future__ import annotations + +import random +from pathlib import Path + +import numpy as np +import torch + +from ai_mouse import generate + +CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [ + ((100, 200), (900, 400)), # horizontal 800px + ((500, 500), (500, 100)), # vertical 400px upward + ((200, 600), (800, 200)), # 720px diagonal + ((100, 100), (130, 110)), # very short 31px + ((50, 50), (1500, 900)), # very long 1700px + ((400, 300), (500, 300)), # short horizontal 100px + ((300, 300), (700, 700)), # 45° diagonal + ((600, 400), (200, 100)), # reverse diagonal +] +SEEDS = (0, 1, 2, 3) + + +def main() -> None: + out: dict[str, np.ndarray] = {} + for case_idx, (start, end) in enumerate(CASES): + for seed in SEEDS: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + pts = generate(start=start, end=end) + out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64) + out_path = Path("tests/unit/data/golden_mouse.npz") + np.savez_compressed(out_path, **out) + print(f"Wrote {len(out)} golden traces to {out_path}") + + +if __name__ == "__main__": + main() +``` + +- [ ] **Step 3: Run the script** + +```bash +uv run python scripts/build_golden_mouse.py +``` + +Expected output: +``` +Wrote 32 golden traces to tests/unit/data/golden_mouse.npz +``` + +- [ ] **Step 4: Inspect the npz** + +```bash +uv run python -c " +import numpy as np +z = np.load('tests/unit/data/golden_mouse.npz') +print('keys:', list(z.keys())[:4], '...') +print('case0_seed0 shape:', z['case0_seed0'].shape) +print('case0_seed0 first 3 rows:', z['case0_seed0'][:3]) +print('case0_seed0 last 2 rows (clicks):', z['case0_seed0'][-2:]) +" +``` + +Expected: 32 keys, each shape (66, 3) — 64 moves + 2 click events. Last two rows share x,y; t increments. + +- [ ] **Step 5: Commit the golden file (not the script yet)** + +```bash +git add tests/unit/data/golden_mouse.npz scripts/build_golden_mouse.py +git commit -m "test: capture mouse generate() golden output (pre-migration)" +``` + +--- + +### Task 0.3: Build scroll golden npz + +**Files:** +- Create: `scripts/build_golden_scroll.py` (temporary) +- Output: `tests/unit/data/golden_scroll.npz` + +- [ ] **Step 1: Create the script** + +Create `scripts/build_golden_scroll.py`: + +```python +"""Capture golden scroll event sequences from current torch implementation.""" +from __future__ import annotations + +import random +from pathlib import Path + +import numpy as np +import torch + +from ai_mouse import generate_scroll + +CASES: list[tuple[int, int, str]] = [ + (0, 1500, "target"), + (0, 500, "precise"), + (0, 5000, "fast"), + (2000, 0, "target"), # upward + (0, 800, "precise"), + (0, 3500, "fast"), + (1000, 1200, "precise"), # tiny scroll + (0, 10000, "fast"), # very long +] +SEEDS = (0, 1, 2, 3) + + +def main() -> None: + out: dict[str, np.ndarray] = {} + for case_idx, (start_y, end_y, mode) in enumerate(CASES): + for seed in SEEDS: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + events = generate_scroll(start_y, end_y, mode=mode) + arr = np.array( + [[e["deltaY"], e["deltaMode"], e["t"]] for e in events], + dtype=np.int64, + ) + out[f"case{case_idx}_seed{seed}"] = arr + out_path = Path("tests/unit/data/golden_scroll.npz") + np.savez_compressed(out_path, **out) + print(f"Wrote {len(out)} scroll golden traces to {out_path}") + + +if __name__ == "__main__": + main() +``` + +- [ ] **Step 2: Run it** + +```bash +uv run python scripts/build_golden_scroll.py +``` + +Expected: `Wrote 32 scroll golden traces to tests/unit/data/golden_scroll.npz` + +- [ ] **Step 3: Commit** + +```bash +git add tests/unit/data/golden_scroll.npz scripts/build_golden_scroll.py +git commit -m "test: capture scroll generate() golden output (pre-migration)" +``` + +--- + +## Phase 1: Move dev code out of the `ai_mouse/` package + +After Phase 1, `ai_mouse/` package contains ONLY inference-related modules (still torch-based for now). All training/server/collector code lives under `tools/`. The library API `from ai_mouse import generate` still works because we haven't touched it yet. + +### Task 1.1: Scaffold `tools/` directory + +**Files:** +- Create: `tools/__init__.py` + +- [ ] **Step 1: Create tools/ and an empty __init__.py** + +```bash +mkdir -p tools/scroll +touch tools/__init__.py tools/scroll/__init__.py +``` + +- [ ] **Step 2: Verify** + +```bash +ls tools/ tools/scroll/ +``` + +Expected: `__init__.py` in both. + +- [ ] **Step 3: Commit** + +```bash +git add tools/__init__.py tools/scroll/__init__.py +git commit -m "chore: scaffold tools/ directory" +``` + +--- + +### Task 1.2: Move trainer + models + utils + config to tools/ + +Move the torch-using mouse modules together so internal imports stay consistent within one commit. + +**Files:** +- Move: `ai_mouse/trainer.py` → `tools/trainer.py` +- Move: `ai_mouse/models.py` → `tools/models.py` +- Move: `ai_mouse/utils.py` → `tools/utils.py` +- Move: `ai_mouse/config.py` → `tools/config.py` +- Modify: `ai_mouse/generator.py` (update imports) + +- [ ] **Step 1: git mv the files** + +```bash +git mv ai_mouse/trainer.py tools/trainer.py +git mv ai_mouse/models.py tools/models.py +git mv ai_mouse/utils.py tools/utils.py +git mv ai_mouse/config.py tools/config.py +``` + +- [ ] **Step 2: Update imports inside moved files** + +In `tools/trainer.py`, replace: +- `from ai_mouse.config import TrainConfig` → `from tools.config import TrainConfig` +- `from ai_mouse.coord import encode_trajectory` → `from ai_mouse.coord import encode_trajectory` (unchanged — coord stays in package) +- `from ai_mouse.models import TrajectoryFlowModel` → `from tools.models import TrajectoryFlowModel` +- `from ai_mouse.utils import resample_arc` → `from tools.utils import resample_arc` + +In `tools/utils.py`: no imports to change (pure numpy). +In `tools/models.py`: no imports to change (torch only). +In `tools/config.py`: no imports to change. + +- [ ] **Step 3: Update `ai_mouse/generator.py` to import torch model from tools** + +Find the imports near the top: + +```python +from ai_mouse.config import GenerateConfig +from ai_mouse.coord import decode_trajectory +from ai_mouse.models import TrajectoryFlowModel +from ai_mouse.utils import resample_arc +``` + +Replace with: + +```python +from ai_mouse.coord import decode_trajectory +from tools.config import GenerateConfig +from tools.models import TrajectoryFlowModel +from tools.utils import resample_arc +``` + +(Note: `GenerateConfig` also moved with `config.py`. We'll yank this cross-boundary import in Phase 4 when generator.py is replaced entirely.) + +- [ ] **Step 4: Verify package imports** + +```bash +uv run python -c "from ai_mouse import generate; print(generate.__module__)" +``` + +Expected: prints `ai_mouse.generator` with no ImportError. + +- [ ] **Step 5: Run existing tests** + +```bash +uv run pytest tests/test_generator.py tests/test_trainer.py tests/test_models.py -v +``` + +Expected: all pass (some test files may need import updates — see next step). + +- [ ] **Step 6: Update test imports if needed** + +In `tests/test_trainer.py`, `tests/test_models.py`, `tests/conftest.py`, replace: +- `from ai_mouse.trainer import ...` → `from tools.trainer import ...` +- `from ai_mouse.models import ...` → `from tools.models import ...` +- `from ai_mouse.config import TrainConfig` → `from tools.config import TrainConfig` + +- [ ] **Step 7: Re-run tests** + +```bash +uv run pytest tests/test_generator.py tests/test_trainer.py tests/test_models.py -v +``` + +Expected: all pass. + +- [ ] **Step 8: Commit** + +```bash +git add -A +git commit -m "refactor: move trainer/models/utils/config to tools/" +``` + +--- + +### Task 1.3: Move scroll trainer / models / collector + +**Files:** +- Move: `ai_mouse/scroll/trainer.py` → `tools/scroll/trainer.py` +- Move: `ai_mouse/scroll/models.py` → `tools/scroll/models.py` +- Move: `ai_mouse/scroll/collector.py` → `tools/scroll/collector.py` +- Modify: `ai_mouse/scroll/__init__.py`, `ai_mouse/scroll/generator.py` + +- [ ] **Step 1: git mv** + +```bash +git mv ai_mouse/scroll/trainer.py tools/scroll/trainer.py +git mv ai_mouse/scroll/models.py tools/scroll/models.py +git mv ai_mouse/scroll/collector.py tools/scroll/collector.py +``` + +- [ ] **Step 2: Update imports inside moved files** + +In `tools/scroll/trainer.py`: +- `from ai_mouse.scroll.models import ScrollCVAE` → `from tools.scroll.models import ScrollCVAE` +- `from ai_mouse.config import ScrollTrainConfig` → `from tools.config import ScrollTrainConfig` + +In `tools/scroll/collector.py`: +- `from ai_mouse.config import SCROLL_MODES, ScrollModeConfig` → `from tools.config import SCROLL_MODES, ScrollModeConfig` + +- [ ] **Step 3: Update `ai_mouse/scroll/generator.py`** + +Replace `from ai_mouse.scroll.models import ScrollCVAE` with `from tools.scroll.models import ScrollCVAE`. + +- [ ] **Step 4: Strip stale imports from `ai_mouse/scroll/__init__.py`** + +Read current content first: + +```bash +cat ai_mouse/scroll/__init__.py +``` + +Edit it to only re-export `generate_scroll` (the only public surface that stays in the package): + +```python +"""Scroll wheel event generation (inference only).""" +from ai_mouse.scroll.generator import generate_scroll + +__all__ = ["generate_scroll"] +``` + +- [ ] **Step 5: Update test imports** + +In `tests/test_scroll_trainer.py`, `tests/test_scroll_models.py`, `tests/test_scroll_collector.py`: +- `from ai_mouse.scroll.trainer import ...` → `from tools.scroll.trainer import ...` +- `from ai_mouse.scroll.models import ...` → `from tools.scroll.models import ...` +- `from ai_mouse.scroll.collector import ...` → `from tools.scroll.collector import ...` + +In `tests/conftest.py`: +- `from ai_mouse.scroll.models import ScrollCVAE` → `from tools.scroll.models import ScrollCVAE` + +- [ ] **Step 6: Run scroll tests** + +```bash +uv run pytest tests/test_scroll_*.py -v +``` + +Expected: all pass. + +- [ ] **Step 7: Commit** + +```bash +git add -A +git commit -m "refactor(scroll): move trainer/models/collector to tools/scroll/" +``` + +--- + +### Task 1.4: Move mouse collector + +**Files:** +- Move: `ai_mouse/collector.py` → `tools/collector.py` + +- [ ] **Step 1: git mv + import fix** + +```bash +git mv ai_mouse/collector.py tools/collector.py +``` + +In `tools/collector.py`, replace any `from ai_mouse.config import ...` with `from tools.config import ...`. + +- [ ] **Step 2: Search for callers** + +```bash +grep -rn "from ai_mouse.collector" --include="*.py" +grep -rn "from ai_mouse import collector" --include="*.py" +``` + +Update each hit to `from tools.collector import ...`. + +- [ ] **Step 3: Run tests touching collector** + +```bash +uv run pytest tests/ -k collector -v +``` + +Expected: pass. + +- [ ] **Step 4: Commit** + +```bash +git add -A +git commit -m "refactor: move collector to tools/" +``` + +--- + +### Task 1.5: Move server/ + +**Files:** +- Move: `ai_mouse/server/` → `tools/server/` +- Modify: `tools/server/__init__.py` (path resolution to static/) + +- [ ] **Step 1: git mv** + +```bash +git mv ai_mouse/server tools/server +``` + +- [ ] **Step 2: Fix imports inside tools/server/** + +For each file in `tools/server/` (`__init__.py`, `deps.py`, `routes_collect.py`, `routes_train.py`, `routes_verify.py`, `routes_scroll.py`), replace: +- `from ai_mouse.collector import Collector` → `from tools.collector import Collector` +- `from ai_mouse.scroll.collector import ScrollCollector` → `from tools.scroll.collector import ScrollCollector` +- `from ai_mouse.scroll.trainer import train as train_scroll` → `from tools.scroll.trainer import train as train_scroll` +- `from ai_mouse.trainer import train` → `from tools.trainer import train` +- `from ai_mouse.config import ...` → `from tools.config import ...` + +Keep these unchanged (they're library API): +- `from ai_mouse import generate, generate_scroll` + +- [ ] **Step 3: Fix static path resolution in `tools/server/__init__.py`** + +The current code reads: + +```python +_HERE = Path(__file__).resolve().parent +_STATIC_DIR = _HERE.parent.parent / "static" +``` + +After moving, `_HERE` is `tools/server/` so `.parent.parent` becomes the repo root — already correct. Verify by: + +```bash +uv run python -c "from tools.server import create_app; app = create_app(); print('app routes:', len(app.routes))" +``` + +Expected: no error; prints route count. + +- [ ] **Step 4: Update test imports** + +In `tests/test_server.py`: +- `from ai_mouse.server import create_app` → `from tools.server import create_app` +- Any `from ai_mouse.server.X import Y` → `from tools.server.X import Y` +- `import ai_mouse.server.deps as deps_module` (if present) → `import tools.server.deps as deps_module` + +- [ ] **Step 5: Run server tests** + +```bash +uv run pytest tests/test_server.py -v +``` + +Expected: pass. + +- [ ] **Step 6: Commit** + +```bash +git add -A +git commit -m "refactor: move server/ to tools/server/" +``` + +--- + +### Task 1.6: Move eval/ and data_adapters/ + +**Files:** +- Move: `ai_mouse/eval/` → `tools/eval/` +- Move: `ai_mouse/data_adapters/` → `tools/data_adapters/` + +- [ ] **Step 1: git mv** + +```bash +git mv ai_mouse/eval tools/eval +git mv ai_mouse/data_adapters tools/data_adapters +``` + +- [ ] **Step 2: Fix imports in moved files** + +In `tools/eval/__main__.py`: +- `from ai_mouse.eval.report import build_report` → `from tools.eval.report import build_report` + +In `tools/eval/report.py`: +- `from ai_mouse.eval.metrics import ...` → `from tools.eval.metrics import ...` + +In `tools/data_adapters/__main__.py`: +- `from ai_mouse.data_adapters.balabit import main` → `from tools.data_adapters.balabit import main` + +In `tools/data_adapters/balabit.py`: +- `from ai_mouse.config import BalabitAdapterConfig` → `from tools.config import BalabitAdapterConfig` + +- [ ] **Step 3: Update test imports** + +In `tests/test_eval_metrics.py`: +- `from ai_mouse.eval.metrics import ...` → `from tools.eval.metrics import ...` + +In `tests/test_balabit_adapter.py`: +- `from ai_mouse.data_adapters.balabit import ...` → `from tools.data_adapters.balabit import ...` + +- [ ] **Step 4: Run tests** + +```bash +uv run pytest tests/test_eval_metrics.py tests/test_balabit_adapter.py -v +``` + +Expected: pass. + +- [ ] **Step 5: Commit** + +```bash +git add -A +git commit -m "refactor: move eval/ and data_adapters/ to tools/" +``` + +--- + +### Task 1.7: Move CLI dispatcher + +**Files:** +- Move: `ai_mouse/__main__.py` → `tools/__main__.py` +- Modify: `tools/__main__.py` (update internal subcommand wiring) + +- [ ] **Step 1: git mv** + +```bash +git mv ai_mouse/__main__.py tools/__main__.py +``` + +- [ ] **Step 2: Update internal imports** + +In `tools/__main__.py`: +- `from ai_mouse.trainer import train` → `from tools.trainer import train` +- `from ai_mouse.eval.__main__ import main as eval_main` → `from tools.eval.__main__ import main as eval_main` +- `from ai_mouse.data_adapters.balabit import main as bal_main` → `from tools.data_adapters.balabit import main as bal_main` + +- [ ] **Step 3: Verify CLI dispatch** + +```bash +uv run python -m tools --help +``` + +Expected: prints help showing `train`, `eval`, `balabit-adapter` subcommands. + +```bash +uv run python -m tools train --help +``` + +Expected: prints `train`-specific args. + +- [ ] **Step 4: Commit** + +```bash +git add -A +git commit -m "refactor: move CLI dispatcher to tools/__main__.py" +``` + +--- + +### Task 1.8: Convert root `main.py` to `tools/serve.py` + +**Files:** +- Move: `main.py` → `tools/serve.py` + +- [ ] **Step 1: git mv** + +```bash +git mv main.py tools/serve.py +``` + +- [ ] **Step 2: Fix imports in tools/serve.py** + +```python +from tools.server import create_app +``` + +(was `from ai_mouse.server import create_app`) + +- [ ] **Step 3: Verify it starts** + +In one terminal: `uv run python tools/serve.py`. In another: `curl http://127.0.0.1:8765/api/status` (or similar status endpoint). Kill the server. Then: + +```bash +uv run python -c "from tools.serve import app; print('app:', app)" +``` + +Expected: prints `app: ` without error. + +- [ ] **Step 4: Commit** + +```bash +git add -A +git commit -m "refactor: move web entry main.py to tools/serve.py" +``` + +--- + +### Task 1.9: Split tests into `tests/unit/` and `tests/tools/` + +**Files:** +- Move test files based on dependency: + - `tests/unit/`: `test_coord.py`, `test_generator.py` (still uses torch via current `generate()` — KEEP in unit; will be rewritten in Phase 4) + - `tests/tools/`: `test_trainer.py`, `test_models.py`, `test_server.py`, `test_scroll_*.py`, `test_eval_metrics.py`, `test_balabit_adapter.py` + +Special case: `test_generator.py` and `test_coord.py` test the library API — they belong in `tests/unit/`. They depend on torch transitively today (via the current generator.py) but in Phase 4 they will not. Move them now to `tests/unit/`; they will keep working through both phases. + +- [ ] **Step 1: Create test dirs and split** + +```bash +mkdir -p tests/unit tests/tools +git mv tests/test_coord.py tests/unit/test_coord.py +git mv tests/test_generator.py tests/unit/test_generator.py +git mv tests/test_trainer.py tests/tools/test_trainer.py +git mv tests/test_models.py tests/tools/test_models.py +git mv tests/test_server.py tests/tools/test_server.py +git mv tests/test_scroll_collector.py tests/tools/test_scroll_collector.py +git mv tests/test_scroll_generator.py tests/unit/test_scroll_generator.py +git mv tests/test_scroll_models.py tests/tools/test_scroll_models.py +git mv tests/test_scroll_trainer.py tests/tools/test_scroll_trainer.py +git mv tests/test_eval_metrics.py tests/tools/test_eval_metrics.py +git mv tests/test_balabit_adapter.py tests/tools/test_balabit_adapter.py +``` + +- [ ] **Step 2: Split conftest.py** + +Current `tests/conftest.py` provides `model_dir` and `scroll_model_dir` fixtures that use torch. These are used by tests that will end up in `tests/tools/` (the torch-using ones). Move them there: + +```bash +git mv tests/conftest.py tests/tools/conftest.py +``` + +Create empty `tests/unit/conftest.py`: + +```python +"""Fixtures for library-only tests (no torch).""" +``` + +- [ ] **Step 3: Add __init__.py if pytest needs them** + +```bash +touch tests/unit/__init__.py tests/tools/__init__.py +``` + +(`tests/__init__.py` already exists.) + +- [ ] **Step 4: Run both directories separately** + +```bash +uv run pytest tests/unit -v +uv run pytest tests/tools -v +``` + +Expected: both pass. Some tests in tests/unit may still touch torch indirectly via the current generator.py — that's OK, will be cleared in Phase 4. + +- [ ] **Step 5: Commit** + +```bash +git add -A +git commit -m "refactor(tests): split into tests/unit and tests/tools" +``` + +--- + +### Task 1.10: Verify whole Phase 1 outcome + +- [ ] **Step 1: Inspect package surface** + +```bash +ls ai_mouse/ +``` + +Expected (Phase 1 end state): `__init__.py`, `coord.py`, `generator.py`, `scroll/` (with `__init__.py`, `generator.py` only). + +`ai_mouse/scroll/`: + +```bash +ls ai_mouse/scroll/ +``` + +Expected: `__init__.py`, `generator.py`. + +- [ ] **Step 2: Verify imports from each side still work** + +```bash +uv run python -c " +from ai_mouse import generate, generate_scroll +print('Library import OK') +from tools.trainer import train +from tools.scroll.trainer import train as st +from tools.server import create_app +from tools.eval.metrics import compute_speed +print('Tools imports OK') +" +``` + +Expected: prints both OK lines. + +- [ ] **Step 3: Run full test suite** + +```bash +uv run pytest tests/ -v +``` + +Expected: all green. + +--- + +## Phase 2: Switch to `src/` layout + tighten pyproject + +### Task 2.1: git mv `ai_mouse` → `src/ai_mouse` + +**Files:** +- Move: `ai_mouse/` → `src/ai_mouse/` + +- [ ] **Step 1: Move the package** + +```bash +mkdir -p src +git mv ai_mouse src/ai_mouse +``` + +- [ ] **Step 2: Verify nothing inside needs path updates** + +The package code uses absolute imports like `from ai_mouse.coord import ...`. After the move, `ai_mouse` is still importable (because src/ becomes a path entry for setuptools/hatchling). Sanity-check there are no hard-coded paths in the source: + +```bash +grep -rn "ai_mouse/" src/ai_mouse/ --include="*.py" +``` + +Expected: only string matches inside docstrings/comments, no live `Path("ai_mouse/...")` constructions. + +- [ ] **Step 3: Commit** + +```bash +git add -A +git commit -m "refactor: switch to src/ layout" +``` + +--- + +### Task 2.2: Rewrite pyproject.toml (hatchling + tightened deps) + +**Files:** +- Modify: `pyproject.toml` + +- [ ] **Step 1: Backup current pyproject** + +```bash +cp pyproject.toml pyproject.toml.bak +``` + +- [ ] **Step 2: Write the new pyproject.toml** + +Replace the entire file with: + +```toml +[project] +name = "ai-mouse" +version = "0.2.0" +description = "Human-like mouse trajectory and scroll wheel event generator (ONNX Runtime SDK)." +requires-python = ">=3.12,<3.14" +dependencies = [ + "numpy>=1.26.0", + "onnxruntime>=1.17.0", +] + +[project.urls] +Repository = "https://github.com//ai_mouse" + +[dependency-groups] +dev = [ + "torch>=2.2.0", + "fastapi>=0.111.0", + "uvicorn>=0.29.0", + "scipy>=1.10.0", + "matplotlib>=3.8.0", + "pytest>=8.0.0", + "pytest-asyncio>=0.23.0", + "httpx>=0.27.0", + "onnx>=1.15.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/ai_mouse"] + +[tool.hatch.build.targets.wheel.force-include] +"src/ai_mouse/assets" = "ai_mouse/assets" + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] +``` + +Notes: +- `onnx` is in `[dev]` only because the export script in `tools/export_onnx.py` uses it; runtime doesn't. +- The `force-include` line is harmless before assets/ exists; it becomes load-bearing in Phase 3. + +- [ ] **Step 3: Re-sync dev environment** + +```bash +uv sync --group dev +``` + +Expected: completes without error. `uv.lock` updates. + +- [ ] **Step 4: Run all tests** + +```bash +uv run pytest tests/ -v +``` + +Expected: all pass. + +- [ ] **Step 5: Delete the backup** + +```bash +rm pyproject.toml.bak +``` + +- [ ] **Step 6: Commit** + +```bash +git add pyproject.toml uv.lock +git commit -m "build: switch to hatchling + src layout; tighten runtime deps" +``` + +--- + +### Task 2.3: Smoke-test wheel build + +**Files:** +- (no files modified, just verify build) + +- [ ] **Step 1: Build the wheel** + +```bash +uv build +``` + +Expected: produces `dist/ai_mouse-0.2.0-py3-none-any.whl` and `dist/ai_mouse-0.2.0.tar.gz`. + +- [ ] **Step 2: Inspect wheel contents** + +```bash +uv run python -c " +import zipfile +with zipfile.ZipFile('dist/ai_mouse-0.2.0-py3-none-any.whl') as z: + for n in z.namelist(): + print(n) +" +``` + +Expected: shows `ai_mouse/__init__.py`, `ai_mouse/generator.py`, `ai_mouse/coord.py`, etc. No `tools/` content; no `tests/`. + +- [ ] **Step 3: Try installing into a clean venv** + +```bash +uv venv .venv-clean +.venv-clean/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl +.venv-clean/Scripts/python -c "import ai_mouse; print(ai_mouse.__file__)" +``` + +Expected: import works. Note `torch` is NOT installed in this venv, so `from ai_mouse import generate` will FAIL right now (current generator.py still imports torch). That's expected pre-Phase-4 — just confirm `import ai_mouse` itself succeeds (it doesn't trigger generator.py). + +Actually `ai_mouse/__init__.py` does `from ai_mouse.generator import generate`, which transitively imports torch. So this import WILL fail. Expected outcome: + +``` +ModuleNotFoundError: No module named 'torch' +``` + +Confirms the wheel content is correct but the runtime promise isn't met yet — exactly the state we expect at Phase 2 end. Document this in commit message. + +- [ ] **Step 4: Clean up** + +```bash +rm -rf .venv-clean dist/ +``` + +- [ ] **Step 5: No commit needed (verification only)** + +--- + +## Phase 3: ONNX export + +### Task 3.1: Write the mouse-model export portion of `tools/export_onnx.py` + +**Files:** +- Create: `tools/export_onnx.py` + +- [ ] **Step 1: Create the file with imports and helpers** + +Create `tools/export_onnx.py`: + +```python +"""Export trained PyTorch checkpoints to ONNX for the inference SDK. + +Usage: + uv run python tools/export_onnx.py \ + --flow-ckpt data/models_v2 \ + --scroll-ckpt data/scroll_models \ + --output src/ai_mouse/assets/ + +Produces: + /flow_model.onnx + /scroll_decoder.onnx + /click_dist.json + /duration_dist.json + /train_config.json + /scroll_config.json + +A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails +the .onnx files are deleted to prevent shipping broken weights. +""" +from __future__ import annotations + +import argparse +import json +import logging +import shutil +import sys +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + +_ATOL = 1e-4 +``` + +- [ ] **Step 2: Add `export_flow_model` function** + +Append to `tools/export_onnx.py`: + +```python +def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path: + """Export TrajectoryFlowModel to ONNX. + + Args: + ckpt_dir: directory with flow_model.pt and train_config.json. + out_dir: destination directory (created if missing). + + Returns: + Path to the written flow_model.onnx. + """ + from tools.models import TrajectoryFlowModel + + config_path = ckpt_dir / "train_config.json" + cfg = json.loads(config_path.read_text()) + seq_len = int(cfg["seq_len"]) + d_model = int(cfg["d_model"]) + nhead = int(cfg["nhead"]) + num_layers = int(cfg["num_layers"]) + dim_feedforward = int(cfg["dim_feedforward"]) + cond_dim = int(cfg.get("cond_dim", 3)) + + model = TrajectoryFlowModel( + seq_len=seq_len, + d_model=d_model, + nhead=nhead, + num_layers=num_layers, + dim_feedforward=dim_feedforward, + cond_dim=cond_dim, + dropout=0.0, # disable dropout for export + ) + state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True) + model.load_state_dict(state) + model.eval() + + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "flow_model.onnx" + + dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32) + dummy_t = torch.zeros(1, dtype=torch.float32) + dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32) + + torch.onnx.export( + model, + (dummy_x, dummy_t, dummy_cond), + str(out_path), + input_names=["x_t", "t", "cond"], + output_names=["v"], + dynamic_axes={ + "x_t": {0: "batch"}, + "t": {0: "batch"}, + "cond": {0: "batch"}, + "v": {0: "batch"}, + }, + opset_version=17, + do_constant_folding=True, + ) + logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6) + return out_path +``` + +- [ ] **Step 3: Add a quick sanity test (manual run)** + +In a python shell: + +```bash +uv run python -c " +from pathlib import Path +from tools.export_onnx import export_flow_model +out = export_flow_model(Path('data/models_v2'), Path('/tmp/test_export')) +print('Wrote:', out) +" +``` + +Expected: prints the output path and a size like 2-3 MB. The file exists. + +- [ ] **Step 4: Commit** + +```bash +git add tools/export_onnx.py +git commit -m "feat(tools): add export_flow_model for ONNX export" +``` + +--- + +### Task 3.2: Add scroll-decoder export + +**Files:** +- Modify: `tools/export_onnx.py` + +- [ ] **Step 1: Define ScrollDecoder wrapper module** + +Append to `tools/export_onnx.py`: + +```python +class _ScrollDecoder(torch.nn.Module): + """Wraps ScrollCVAE.decode for ONNX export. + + The full ScrollCVAE is encoder+decoder; inference only needs decoder. + """ + + def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int): + super().__init__() + self.dec_h0 = dec_h0 + self.dec_gru = dec_gru + self.dec_out = dec_out + self.seq_len = seq_len + self.hidden = hidden + + def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + b = z.shape[0] + zc = torch.cat([z, cond], dim=-1) + h0_flat = self.dec_h0(zc) + h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous() + inp = zc.unsqueeze(1).expand(b, self.seq_len, -1) + out, _ = self.dec_gru(inp, h0) + return self.dec_out(out) +``` + +- [ ] **Step 2: Add `export_scroll_decoder` function** + +Append: + +```python +def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path: + """Export ScrollCVAE decoder to ONNX.""" + from tools.scroll.models import ScrollCVAE + + config_path = ckpt_dir / "scroll_config.json" + cfg = json.loads(config_path.read_text()) + seq_len = int(cfg["seq_len"]) + latent_dim = int(cfg["latent_dim"]) + hidden = int(cfg["hidden"]) + cond_dim = int(cfg["cond_dim"]) + + full = ScrollCVAE( + seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim + ) + state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True) + full.load_state_dict(state) + full.eval() + + decoder = _ScrollDecoder( + dec_h0=full.dec_h0, + dec_gru=full.dec_gru, + dec_out=full.dec_out, + seq_len=seq_len, + hidden=hidden, + ) + decoder.eval() + + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "scroll_decoder.onnx" + + dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32) + dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32) + + torch.onnx.export( + decoder, + (dummy_z, dummy_cond), + str(out_path), + input_names=["z", "cond"], + output_names=["seq"], + dynamic_axes={ + "z": {0: "batch"}, + "cond": {0: "batch"}, + "seq": {0: "batch"}, + }, + opset_version=17, + do_constant_folding=True, + ) + logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3) + return out_path +``` + +- [ ] **Step 3: Manual sanity test** + +```bash +uv run python -c " +from pathlib import Path +from tools.export_onnx import export_scroll_decoder +out = export_scroll_decoder(Path('data/scroll_models'), Path('/tmp/test_export')) +print('Wrote:', out) +" +``` + +Expected: prints path; file <300 KB. + +- [ ] **Step 4: Commit** + +```bash +git add tools/export_onnx.py +git commit -m "feat(tools): add export_scroll_decoder for ONNX export" +``` + +--- + +### Task 3.3: Add PyTorch vs ORT parity check + +**Files:** +- Modify: `tools/export_onnx.py` + +- [ ] **Step 1: Add parity helpers** + +Append to `tools/export_onnx.py`: + +```python +def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None: + """Verify ONNX flow model matches PyTorch output on random input.""" + import onnxruntime as ort + from tools.models import TrajectoryFlowModel + + cfg = json.loads((ckpt_dir / "train_config.json").read_text()) + seq_len = int(cfg["seq_len"]) + cond_dim = int(cfg.get("cond_dim", 3)) + + model = TrajectoryFlowModel( + seq_len=seq_len, + d_model=int(cfg["d_model"]), + nhead=int(cfg["nhead"]), + num_layers=int(cfg["num_layers"]), + dim_feedforward=int(cfg["dim_feedforward"]), + cond_dim=cond_dim, + dropout=0.0, + ) + model.load_state_dict( + torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True) + ) + model.eval() + + torch.manual_seed(42) + np.random.seed(42) + x = torch.randn(2, seq_len, 3, dtype=torch.float32) + t = torch.tensor([0.0, 0.5], dtype=torch.float32) + cond = torch.randn(2, cond_dim, dtype=torch.float32) + + with torch.no_grad(): + torch_out = model(x, t, cond).numpy() + + sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_out = sess.run( + ["v"], + { + "x_t": x.numpy(), + "t": t.numpy(), + "cond": cond.numpy(), + }, + )[0] + + if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3): + max_diff = float(np.abs(torch_out - ort_out).max()) + raise RuntimeError( + f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}" + ) + logger.info("Flow model parity OK (atol=%.0e)", _ATOL) + + +def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None: + """Verify ONNX scroll decoder matches PyTorch decoder output.""" + import onnxruntime as ort + from tools.scroll.models import ScrollCVAE + + cfg = json.loads((ckpt_dir / "scroll_config.json").read_text()) + seq_len = int(cfg["seq_len"]) + latent_dim = int(cfg["latent_dim"]) + cond_dim = int(cfg["cond_dim"]) + + full = ScrollCVAE( + seq_len=seq_len, + latent_dim=latent_dim, + hidden=int(cfg["hidden"]), + cond_dim=cond_dim, + ) + full.load_state_dict( + torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True) + ) + full.eval() + + torch.manual_seed(7) + z = torch.randn(2, latent_dim, dtype=torch.float32) + cond = torch.randn(2, cond_dim, dtype=torch.float32) + + with torch.no_grad(): + torch_out = full.decode(z, cond).numpy() + + sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0] + + if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3): + max_diff = float(np.abs(torch_out - ort_out).max()) + raise RuntimeError( + f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}" + ) + logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL) +``` + +- [ ] **Step 2: Manual test the checks** + +```bash +uv run python -c " +from pathlib import Path +from tools.export_onnx import ( + export_flow_model, export_scroll_decoder, + _check_flow_parity, _check_scroll_parity, +) +import logging; logging.basicConfig(level=logging.INFO) +out = Path('/tmp/test_export') +export_flow_model(Path('data/models_v2'), out) +_check_flow_parity(Path('data/models_v2'), out / 'flow_model.onnx') +export_scroll_decoder(Path('data/scroll_models'), out) +_check_scroll_parity(Path('data/scroll_models'), out / 'scroll_decoder.onnx') +" +``` + +Expected: prints two "parity OK" lines, no exceptions. + +- [ ] **Step 3: Commit** + +```bash +git add tools/export_onnx.py +git commit -m "feat(tools): add ORT vs PyTorch parity check for exports" +``` + +--- + +### Task 3.4: Add CLI `main()` to `tools/export_onnx.py` + +**Files:** +- Modify: `tools/export_onnx.py` + +- [ ] **Step 1: Add main() and __main__ guard** + +Append to `tools/export_onnx.py`: + +```python +def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None: + """Copy JSON metadata files alongside the ONNX models.""" + for name in ("click_dist.json", "duration_dist.json", "train_config.json"): + src = flow_dir / name + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / name) + src = scroll_dir / "scroll_config.json" + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / "scroll_config.json") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0]) + p.add_argument("--flow-ckpt", type=Path, required=True) + p.add_argument("--scroll-ckpt", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + args = p.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + args.output.mkdir(parents=True, exist_ok=True) + + flow_onnx = export_flow_model(args.flow_ckpt, args.output) + scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output) + + try: + _check_flow_parity(args.flow_ckpt, flow_onnx) + _check_scroll_parity(args.scroll_ckpt, scroll_onnx) + except RuntimeError as exc: + logger.error("Parity check failed: %s", exc) + flow_onnx.unlink(missing_ok=True) + scroll_onnx.unlink(missing_ok=True) + return 1 + + _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output) + logger.info("Export complete: %s", args.output) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) +``` + +- [ ] **Step 2: Run the full export to produce assets** + +```bash +mkdir -p src/ai_mouse/assets +uv run python tools/export_onnx.py \ + --flow-ckpt data/models_v2 \ + --scroll-ckpt data/scroll_models \ + --output src/ai_mouse/assets/ +``` + +Expected output (timestamps elided): +``` +... INFO Wrote src/ai_mouse/assets/flow_model.onnx (2.x MB) +... INFO Wrote src/ai_mouse/assets/scroll_decoder.onnx (0.x KB) +... INFO Flow model parity OK (atol=1e-04) +... INFO Scroll decoder parity OK (atol=1e-04) +... INFO Export complete: src/ai_mouse/assets +``` + +- [ ] **Step 3: Verify assets directory** + +```bash +ls src/ai_mouse/assets/ +``` + +Expected: `flow_model.onnx`, `scroll_decoder.onnx`, `click_dist.json`, `duration_dist.json`, `train_config.json`, `scroll_config.json`. + +- [ ] **Step 4: Commit assets + script main()** + +```bash +git add tools/export_onnx.py src/ai_mouse/assets/ +git commit -m "feat: export ONNX weights and metadata into src/ai_mouse/assets/" +``` + +--- + +### Task 3.5: Test ONNX export with toy model + +**Files:** +- Create: `tests/tools/test_export_onnx.py` + +- [ ] **Step 1: Write the failing test** + +Create `tests/tools/test_export_onnx.py`: + +```python +"""Validate tools.export_onnx with a tiny synthetic model.""" +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tools.export_onnx import ( + _check_flow_parity, + _check_scroll_parity, + export_flow_model, + export_scroll_decoder, +) + + +@pytest.fixture +def tiny_flow_ckpt(tmp_path: Path) -> Path: + """A flow model with seq_len=8, d_model=16, 1 layer — small but valid.""" + from tools.models import TrajectoryFlowModel + + cfg = { + "seq_len": 8, + "d_model": 16, + "nhead": 2, + "num_layers": 1, + "dim_feedforward": 32, + "cond_dim": 3, + } + model = TrajectoryFlowModel(**cfg, dropout=0.0) + model.eval() + out = tmp_path / "flow_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "flow_model.pt") + (out / "train_config.json").write_text(json.dumps(cfg)) + return out + + +@pytest.fixture +def tiny_scroll_ckpt(tmp_path: Path) -> Path: + """A scroll model with seq_len=4, latent=4, hidden=8.""" + from tools.scroll.models import ScrollCVAE + + cfg = {"seq_len": 4, "latent_dim": 4, "hidden": 8, "cond_dim": 7} + model = ScrollCVAE(**cfg) + model.eval() + out = tmp_path / "scroll_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "scroll_model.pt") + (out / "scroll_config.json").write_text(json.dumps(cfg)) + return out + + +def test_export_flow_model_parity(tiny_flow_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_flow_model(tiny_flow_ckpt, out_dir) + assert onnx_path.exists() + _check_flow_parity(tiny_flow_ckpt, onnx_path) # raises on failure + + +def test_export_scroll_decoder_parity(tiny_scroll_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_scroll_decoder(tiny_scroll_ckpt, out_dir) + assert onnx_path.exists() + _check_scroll_parity(tiny_scroll_ckpt, onnx_path) +``` + +- [ ] **Step 2: Run the tests** + +```bash +uv run pytest tests/tools/test_export_onnx.py -v +``` + +Expected: both pass. + +- [ ] **Step 3: Commit** + +```bash +git add tests/tools/test_export_onnx.py +git commit -m "test(tools): cover export_onnx with tiny synthetic models" +``` + +--- + +## Phase 4: Rewrite library in NumPy + ORT + +### Task 4.1: Create `_coord.py` (private numpy coordinate transforms) + +**Files:** +- Create: `src/ai_mouse/_coord.py` +- Keep (for now): `src/ai_mouse/coord.py` — tools/ still imports it; deleted at end of Phase 4 + +- [ ] **Step 1: Copy coord.py to _coord.py** + +```bash +cp src/ai_mouse/coord.py src/ai_mouse/_coord.py +``` + +- [ ] **Step 2: No content edits needed** + +The file is already pure numpy. Verify: + +```bash +grep -E "^import|^from" src/ai_mouse/_coord.py +``` + +Expected: only `import math` and `import numpy as np`. + +- [ ] **Step 3: Write the test** + +Create `tests/unit/test__coord.py`: + +```python +"""Test the private numpy coordinate transforms.""" +from __future__ import annotations + +import numpy as np + +from ai_mouse._coord import decode_trajectory, encode_trajectory + + +def test_encode_decode_roundtrip() -> None: + points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]]) + start = (100, 200) + end = (500, 300) + encoded = encode_trajectory(points, start, end) + decoded = decode_trajectory(encoded, start, end) + assert np.allclose(decoded, points, atol=1e-6) + + +def test_encode_endpoints() -> None: + """Start should encode to (0,0); end should encode to (1,0).""" + points = np.array([[100.0, 200.0], [500.0, 300.0]]) + encoded = encode_trajectory(points, (100, 200), (500, 300)) + assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6) + assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6) + + +def test_zero_distance_returns_zeros() -> None: + points = np.array([[100.0, 200.0]]) + encoded = encode_trajectory(points, (100, 200), (100, 200)) + assert encoded.shape == (1, 2) + assert np.all(encoded == 0) +``` + +- [ ] **Step 4: Run test** + +```bash +uv run pytest tests/unit/test__coord.py -v +``` + +Expected: 3 pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_coord.py tests/unit/test__coord.py +git commit -m "feat(lib): add private _coord.py with numpy transforms" +``` + +--- + +### Task 4.2: Create `errors.py` + +**Files:** +- Create: `src/ai_mouse/errors.py` + +- [ ] **Step 1: Write the failing test** + +Create `tests/unit/test_errors.py`: + +```python +"""Test the error hierarchy.""" +from __future__ import annotations + +import pytest + +from ai_mouse import errors + + +def test_model_load_error_is_aimouse_error() -> None: + assert issubclass(errors.ModelLoadError, errors.AiMouseError) + + +def test_generation_error_is_aimouse_error() -> None: + assert issubclass(errors.GenerationError, errors.AiMouseError) + + +def test_can_catch_specific_with_general() -> None: + with pytest.raises(errors.AiMouseError): + raise errors.ModelLoadError("test") +``` + +- [ ] **Step 2: Run test, observe failure** + +```bash +uv run pytest tests/unit/test_errors.py -v +``` + +Expected: ImportError on `from ai_mouse import errors`. + +- [ ] **Step 3: Create the module** + +Create `src/ai_mouse/errors.py`: + +```python +"""Exception hierarchy for the ai_mouse library. + +Downstream consumers can catch the umbrella :class:`AiMouseError` +or the specific subclasses for finer control. +""" +from __future__ import annotations + + +class AiMouseError(Exception): + """Base class for all ai_mouse errors.""" + + +class ModelLoadError(AiMouseError): + """Raised when ONNX weights / metadata cannot be loaded.""" + + +class GenerationError(AiMouseError): + """Raised when inference produces an invalid result (e.g. NaN).""" +``` + +- [ ] **Step 4: Run test, observe pass** + +```bash +uv run pytest tests/unit/test_errors.py -v +``` + +Expected: 3 pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/errors.py tests/unit/test_errors.py +git commit -m "feat(lib): add errors module" +``` + +--- + +### Task 4.3: Create `_assets.py` (importlib.resources loader) + +**Files:** +- Create: `src/ai_mouse/_assets.py` + +- [ ] **Step 1: Write the failing test** + +Create `tests/unit/test_assets.py`: + +```python +"""Test the asset path resolver.""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from ai_mouse import _assets +from ai_mouse.errors import ModelLoadError + + +def test_bundled_flow_model_exists() -> None: + p = _assets.bundled_path("flow_model.onnx") + assert p.exists() + assert p.suffix == ".onnx" + + +def test_bundled_train_config_loadable() -> None: + p = _assets.bundled_path("train_config.json") + cfg = json.loads(p.read_text()) + assert "seq_len" in cfg + assert "d_model" in cfg + + +def test_resolve_with_custom_dir(tmp_path: Path) -> None: + (tmp_path / "flow_model.onnx").write_bytes(b"x") + p = _assets.resolve(tmp_path, "flow_model.onnx") + assert p == tmp_path / "flow_model.onnx" + + +def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None: + with pytest.raises(ModelLoadError, match="missing"): + _assets.resolve(tmp_path, "nonexistent.onnx") +``` + +- [ ] **Step 2: Run test, observe failure** + +```bash +uv run pytest tests/unit/test_assets.py -v +``` + +Expected: ImportError on `from ai_mouse import _assets`. + +- [ ] **Step 3: Create the module** + +Create `src/ai_mouse/_assets.py`: + +```python +"""Asset path resolution for bundled ONNX weights and JSON metadata. + +Uses :mod:`importlib.resources` to locate files inside the installed +package, falling back to a user-supplied directory if provided. +""" +from __future__ import annotations + +from importlib.resources import as_file, files +from pathlib import Path + +from ai_mouse.errors import ModelLoadError + +_PACKAGE_ASSETS = "ai_mouse.assets" + + +def bundled_path(name: str) -> Path: + """Return a filesystem path to a bundled asset. + + Args: + name: filename inside the assets/ directory. + + Returns: + A concrete :class:`pathlib.Path`. Note: for zipapp installations + this materialises a temp file; for normal site-packages installs + it points into the package directly. + """ + ref = files(_PACKAGE_ASSETS) / name + # as_file is the canonical way; for non-zip installs this is a no-op + # context that yields the actual path. + with as_file(ref) as p: + # We're inside the with-block; the contextmanager keeps the + # temp file alive only while open. For zip installs we'd need + # to extract to a stable location. For now, all our installs + # are wheel-based (non-zip), so the path is stable after exit. + return Path(p) + + +def resolve(model_path: Path | None, filename: str) -> Path: + """Locate an asset given an optional user-supplied directory. + + Args: + model_path: user-supplied directory, or None to use bundled assets. + filename: file to locate inside the directory. + + Returns: + Absolute path to the asset. + + Raises: + ModelLoadError: if the file does not exist. + """ + if model_path is None: + p = bundled_path(filename) + else: + p = Path(model_path) / filename + if not p.exists(): + raise ModelLoadError(f"Required asset missing: {p}") + return p +``` + +- [ ] **Step 4: Run test, observe pass** + +```bash +uv run pytest tests/unit/test_assets.py -v +``` + +Expected: 4 pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_assets.py tests/unit/test_assets.py +git commit -m "feat(lib): add _assets module for bundled-weight resolution" +``` + +--- + +### Task 4.4: Create `_postprocess.py` skeleton + `gaussian_smooth` + +**Files:** +- Create: `src/ai_mouse/_postprocess.py` +- Create: `tests/unit/test_postprocess.py` + +- [ ] **Step 1: Write failing test** + +Create `tests/unit/test_postprocess.py`: + +```python +"""Tests for trajectory post-processing primitives.""" +from __future__ import annotations + +import numpy as np + +from ai_mouse._postprocess import gaussian_smooth + + +def test_gaussian_smooth_preserves_endpoints() -> None: + x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0]) + result = gaussian_smooth(x, sigma=1.0) + assert result[0] == 1.0 + assert result[-1] == 4.0 + + +def test_gaussian_smooth_short_input_unchanged() -> None: + x = np.array([1.0, 2.0, 3.0]) + result = gaussian_smooth(x, sigma=1.0) + assert np.array_equal(result, x) + + +def test_gaussian_smooth_constant_unchanged() -> None: + x = np.full(20, 7.5) + result = gaussian_smooth(x, sigma=1.0) + assert np.allclose(result, x, atol=1e-6) +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +Expected: ImportError. + +- [ ] **Step 3: Create module + function** + +Create `src/ai_mouse/_postprocess.py`: + +```python +"""Pure-numpy post-processing primitives for trajectory generation. + +All functions are pure (no I/O, no global state) and accept an explicit +:class:`numpy.random.Generator` when randomness is involved. +""" +from __future__ import annotations + +import numpy as np + + +def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray: + """5-tap gaussian smoothing along a 1-D array; endpoints preserved. + + Args: + x: 1-D input array. + sigma: gaussian std. Default 1.0 gives weights ≈ + [0.054, 0.244, 0.403, 0.244, 0.054]. + + Returns: + Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged. + If ``len(x) < 5`` returns a copy of ``x`` (kernel won't fit). + """ + if len(x) < 5: + return x.copy() + kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2) + kernel /= kernel.sum() + padded = np.pad(x, pad_width=2, mode="edge") + smoothed = np.convolve(padded, kernel, mode="valid") + smoothed[0] = x[0] + smoothed[-1] = x[-1] + return smoothed +``` + +- [ ] **Step 4: Run, observe pass** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +Expected: 3 pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py +git commit -m "feat(lib): add gaussian_smooth to _postprocess" +``` + +--- + +### Task 4.5: Add `snap_endpoints` + +**Files:** +- Modify: `src/ai_mouse/_postprocess.py` +- Modify: `tests/unit/test_postprocess.py` + +- [ ] **Step 1: Write failing test (append)** + +Append to `tests/unit/test_postprocess.py`: + +```python +from ai_mouse._postprocess import snap_endpoints + + +def test_snap_endpoints_pins_first_and_last() -> None: + forward = np.linspace(0.1, 0.9, 16) + lateral = np.full(16, 0.5) + f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16) + assert f[0] == 0.0 + assert l[0] == 0.0 + assert f[-1] == 1.0 + assert l[-1] == 0.0 + + +def test_snap_endpoints_preserves_middle() -> None: + forward = np.linspace(0.0, 1.0, 16) + lateral = np.zeros(16) + f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4) + # Points before the last n_snap should be unchanged + assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6) +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_postprocess.py::test_snap_endpoints_pins_first_and_last -v +``` + +Expected: ImportError. + +- [ ] **Step 3: Implement** + +Append to `src/ai_mouse/_postprocess.py`: + +```python +def snap_endpoints( + forward: np.ndarray, + lateral: np.ndarray, + seq_len: int, + n_snap: int = 6, +) -> tuple[np.ndarray, np.ndarray]: + """Force first point to (0,0) and last point to (1,0) with quadratic ease. + + The last ``n_snap`` points are linearly interpolated towards (1, 0) + with quadratic easing, then the first/last points are pinned exactly. + + Args: + forward: (T,) forward coordinates (modified in place). + lateral: (T,) lateral coordinates (modified in place). + seq_len: length of forward/lateral. + n_snap: number of trailing points to ease (capped at seq_len//4). + + Returns: + ``(forward, lateral)`` after modification. + """ + n_snap = min(n_snap, seq_len // 4) + for i in range(n_snap): + alpha = ((i + 1) / n_snap) ** 2 + k = seq_len - n_snap + i + forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha + lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha + forward[0], lateral[0] = 0.0, 0.0 + forward[-1], lateral[-1] = 1.0, 0.0 + return forward, lateral +``` + +- [ ] **Step 4: Run all postprocess tests** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +Expected: all pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py +git commit -m "feat(lib): add snap_endpoints to _postprocess" +``` + +--- + +### Task 4.6: Add `smooth_start`, `enforce_forward_monotonic` + +**Files:** +- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py` + +- [ ] **Step 1: Write tests (append)** + +```python +from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start + + +def test_smooth_start_dampens_lateral() -> None: + forward = np.linspace(0, 1, 16) + lateral = np.full(16, 1.0) + forward[0] = lateral[0] = 0.0 # invariant: snap already done + _, l = smooth_start(forward.copy(), lateral.copy(), n=4) + # Lateral at points 1-4 should be < original (dampened) + assert l[1] < 1.0 + assert l[4] < 1.0 + # Lateral at point 5+ unchanged + assert l[5] == 1.0 + + +def test_enforce_forward_monotonic_repairs_inversions() -> None: + f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0]) + out = enforce_forward_monotonic(f.copy()) + assert np.all(np.diff(out) > 0), out + + +def test_enforce_forward_monotonic_clips_to_unit_interval() -> None: + f = np.array([-0.1, 0.5, 1.2]) + out = enforce_forward_monotonic(f.copy()) + assert out[0] == 0.0 + assert out[-1] == 1.0 +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +- [ ] **Step 3: Implement (append to _postprocess.py)** + +```python +def smooth_start( + forward: np.ndarray, + lateral: np.ndarray, + n: int = 4, +) -> tuple[np.ndarray, np.ndarray]: + """Dampen lateral oscillation in the first ``n`` points. + + Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is + forced non-decreasing locally; lateral is linearly damped towards 0. + """ + n_start_fix = min(n, len(forward) // 4) + for i in range(1, n_start_fix + 1): + blend = i / (n_start_fix + 1) + forward[i] = max(forward[i], forward[i - 1]) + lateral[i] = lateral[i] * blend + return forward, lateral + + +def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray: + """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints.""" + seq_len = len(forward) + for i in range(1, seq_len - 1): + if forward[i] < forward[i - 1]: + forward[i] = forward[i - 1] + 0.001 + forward = np.clip(forward, 0.0, 1.0) + forward[0] = 0.0 + forward[-1] = 1.0 + return forward +``` + +- [ ] **Step 4: Test** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +Expected: all pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py +git commit -m "feat(lib): add smooth_start, enforce_forward_monotonic" +``` + +--- + +### Task 4.7: Add `resample_arc`, `build_timestamps` + +**Files:** +- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py` + +- [ ] **Step 1: Tests** + +Append: + +```python +from ai_mouse._postprocess import build_timestamps, resample_arc + + +def test_resample_arc_identity_when_same_length() -> None: + pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]]) + out = resample_arc(pts, 4) + assert np.allclose(out, pts, atol=1e-6) + + +def test_resample_arc_changes_length() -> None: + pts = np.array([[float(i), 0.0] for i in range(10)]) + out = resample_arc(pts, 5) + assert out.shape == (5, 2) + # Endpoints preserved + assert np.allclose(out[0], pts[0]) + assert np.allclose(out[-1], pts[-1]) + + +def test_build_timestamps_strictly_increasing() -> None: + log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0]) + ts = build_timestamps(log_dt, total_duration_ms=200.0) + assert ts[0] == 0 + assert np.all(np.diff(ts) >= 1) # at least 1 ms apart + + +def test_build_timestamps_total_close_to_target() -> None: + log_dt = np.array([1.0] * 10) + ts = build_timestamps(log_dt, total_duration_ms=300.0) + # Last timestamp should be roughly total - one slot + assert abs(ts[-1] - 270) < 60 # tolerant of clipping +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_postprocess.py::test_resample_arc_identity_when_same_length -v +``` + +- [ ] **Step 3: Implement** + +Append to `_postprocess.py`: + +```python +def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray: + """Resample a 2-D polyline to ``n_points`` along cumulative arc length.""" + arc = np.concatenate( + [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))] + ) + s_new = np.linspace(0, arc[-1], n_points) + return np.stack( + [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])], + axis=1, + ) + + +def build_timestamps( + log_dt: np.ndarray, + total_duration_ms: float, + dt_clip: tuple[float, float] = (2.0, 150.0), +) -> np.ndarray: + """Convert per-step log_dt + total duration to cumulative ms timestamps. + + Args: + log_dt: (N,) array of natural-log step intervals. + total_duration_ms: target total span. The output is scaled so the + sum approximately matches this (modulo dt_clip). + dt_clip: (min, max) per-step clamp in milliseconds. + + Returns: + (N,) integer-rounded cumulative timestamps starting at 0, + strictly increasing. + """ + n = len(log_dt) + dt_raw = np.clip(np.exp(log_dt), 0.0, None) + dt_sum = dt_raw.sum() + if dt_sum > 1e-6: + scale = total_duration_ms / dt_sum + else: + scale = total_duration_ms / max(n, 1) + dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1]) + + t_abs = np.cumsum(dt_ms) + t_abs = np.concatenate([[0.0], t_abs[:-1]]) + + for i in range(1, n): + if t_abs[i] <= t_abs[i - 1]: + t_abs[i] = t_abs[i - 1] + 1.0 + return t_abs +``` + +- [ ] **Step 4: Run** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py +git commit -m "feat(lib): add resample_arc, build_timestamps" +``` + +--- + +### Task 4.8: Add `sample_duration` + `truncnorm_sample` + +**Files:** +- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py` + +- [ ] **Step 1: Tests** + +```python +from ai_mouse._postprocess import sample_duration, truncnorm_sample + + +def test_truncnorm_sample_within_bounds() -> None: + rng = np.random.default_rng(0) + samples = [ + truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500) + ] + arr = np.array(samples) + assert arr.min() >= 20.0 + assert arr.max() <= 300.0 + # Mean roughly close to mu + assert abs(arr.mean() - 80.0) < 5.0 + + +def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None: + rng = np.random.default_rng(0) + # mu far outside [low, high] — rejection will fail + v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng) + assert 20.0 <= v <= 30.0 + + +def test_sample_duration_uses_correct_bin() -> None: + dist_dict = { + "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")], + "params": [ + {"mu_log": 4.0, "sigma_log": 0.01}, # bin 0: dist < 50 + {"mu_log": 5.0, "sigma_log": 0.01}, # bin 1: 50 <= dist < 100 + {"mu_log": 6.0, "sigma_log": 0.01}, # bin 2: 100 <= dist < 200 + ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5, + } + rng = np.random.default_rng(0) + v = sample_duration(dist_dict, 150.0, rng) + # exp(6) ~ 403, with tiny sigma we should land near there + assert 350 < v < 460 +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +- [ ] **Step 3: Implement** + +Append to `_postprocess.py`: + +```python +def truncnorm_sample( + mu: float, + sigma: float, + low: float, + high: float, + rng: np.random.Generator, + max_tries: int = 32, +) -> float: + """Sample from N(mu, sigma) truncated to [low, high] via rejection. + + Falls back to clipping if rejection fails ``max_tries`` times. + """ + for _ in range(max_tries): + v = rng.normal(mu, sigma) + if low <= v <= high: + return float(v) + return float(np.clip(rng.normal(mu, sigma), low, high)) + + +def sample_duration( + duration_dist: dict, + dist: float, + rng: np.random.Generator, +) -> float: + """Sample total trajectory duration (ms) for the given pixel distance. + + Uses per-bin log-normal parameters in ``duration_dist``. + """ + bins = duration_dist["bins"] + params = duration_dist["params"] + bin_idx = len(bins) - 1 + for i in range(len(bins) - 1): + if dist < bins[i + 1]: + bin_idx = i + break + bin_idx = min(bin_idx, len(params) - 1) + mu_log = params[bin_idx]["mu_log"] + sigma_log = params[bin_idx]["sigma_log"] + return float(np.exp(rng.normal(mu_log, sigma_log))) +``` + +- [ ] **Step 4: Test** + +```bash +uv run pytest tests/unit/test_postprocess.py -v +``` + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py +git commit -m "feat(lib): add sample_duration, truncnorm_sample (no scipy)" +``` + +--- + +### Task 4.9: Write `mouse.py` (`MouseModel` + `_get_default_mouse_model`) + +**Files:** +- Create: `src/ai_mouse/mouse.py` + +- [ ] **Step 1: Write test scaffolding** + +Create `tests/unit/test_mouse.py`: + +```python +"""Tests for MouseModel and ai_mouse.generate().""" +from __future__ import annotations + +import numpy as np +import pytest + + +def test_mouse_model_init_default() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + assert m._seq_len > 0 + assert m._session is not None + m.close() + + +def test_mouse_model_generate_returns_correct_shape() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400)) + assert len(pts) == 66 # 64 moves + 2 clicks + for x, y, t in pts: + assert isinstance(x, int) + assert isinstance(y, int) + assert isinstance(t, int) + + +def test_mouse_model_click_false_omits_clicks() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400), click=False) + assert len(pts) == 64 + + +def test_mouse_model_seed_reproducibility() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + a = m.generate((100, 200), (900, 400), seed=42) + b = m.generate((100, 200), (900, 400), seed=42) + assert a == b + + +def test_mouse_model_invalid_path_raises_model_load_error() -> None: + from ai_mouse.mouse import MouseModel + from ai_mouse.errors import ModelLoadError + with pytest.raises(ModelLoadError): + MouseModel(model_path="/nonexistent/path/here") +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_mouse.py -v +``` + +Expected: ImportError. + +- [ ] **Step 3: Implement mouse.py** + +Create `src/ai_mouse/mouse.py`: + +```python +"""MouseModel — ONNX Runtime-backed mouse trajectory generation.""" +from __future__ import annotations + +import json +import math +from collections.abc import Sequence +from pathlib import Path +from typing import Optional + +import numpy as np +import onnxruntime as ort + +from ai_mouse._assets import resolve +from ai_mouse._coord import decode_trajectory +from ai_mouse._postprocess import ( + build_timestamps, + enforce_forward_monotonic, + gaussian_smooth, + resample_arc, + sample_duration, + smooth_start, + snap_endpoints, + truncnorm_sample, +) +from ai_mouse.errors import GenerationError, ModelLoadError + +_N_EULER_STEPS = 10 + + +class MouseModel: + """Persistent ONNX Runtime session for mouse trajectory generation. + + Construct once and reuse across calls — the underlying + ``InferenceSession`` is created lazily in ``__init__`` and kept alive + until :meth:`close` is called. + """ + + def __init__( + self, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, + seed: int | None = None, + ) -> None: + path_obj: Optional[Path] = Path(model_path) if model_path is not None else None + + onnx_path = resolve(path_obj, "flow_model.onnx") + cfg_path = resolve(path_obj, "train_config.json") + click_path = resolve(path_obj, "click_dist.json") + dur_path = resolve(path_obj, "duration_dist.json") + + cfg = json.loads(cfg_path.read_text()) + self._seq_len = int(cfg["seq_len"]) + self._cond_dim = int(cfg.get("cond_dim", 3)) + self._click_params = json.loads(click_path.read_text()) + self._duration_dist = json.loads(dur_path.read_text()) + + try: + self._session = ort.InferenceSession( + str(onnx_path), + providers=list(providers) if providers else ["CPUExecutionProvider"], + ) + except Exception as exc: + raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc + + self._default_seed = seed + self._rng = np.random.default_rng(seed) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def generate( + self, + start: tuple[int, int], + end: tuple[int, int], + n_points: int = 64, + speed: float | None = None, + click: bool = True, + seed: int | None = None, + ) -> list[tuple[int, int, int]]: + """Generate a human-like mouse trajectory from ``start`` to ``end``. + + Args: + start: (x, y) starting pixel coordinate. + end: (x, y) target pixel coordinate. + n_points: number of move points (default 64). + speed: optional multiplier; ``speed=2`` halves the duration. + click: append mouse-down and mouse-up events at the end. + seed: per-call seed overriding the instance seed. + + Returns: + List of (x, y, t_ms) tuples. ``t_ms`` is cumulative from 0. + """ + rng = np.random.default_rng(seed) if seed is not None else self._rng + + sx, sy = float(start[0]), float(start[1]) + ex, ey = float(end[0]), float(end[1]) + dist = max(math.hypot(ex - sx, ey - sy), 1.0) + + total_duration = sample_duration(self._duration_dist, dist, rng) + if speed is not None and speed > 0: + total_duration /= speed + total_duration = max(total_duration, 10.0) + + cond = np.array( + [ + dist / 2000.0, + math.log(dist / 100.0), + math.log(total_duration / 500.0), + ], + dtype=np.float32, + )[None] + + # Euler ODE + x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32) + dt = 1.0 / _N_EULER_STEPS + for step in range(_N_EULER_STEPS): + t = np.full((1,), step * dt, dtype=np.float32) + v = self._session.run( + ["v"], {"x_t": x, "t": t, "cond": cond} + )[0] + x = x + v * dt + + if not np.all(np.isfinite(x)): + raise GenerationError("Trajectory contains NaN/Inf after Euler integration") + + forward = x[0, :, 0].copy() + lateral = x[0, :, 1].copy() + log_dt = x[0, :, 2].copy() + + # Spatial post-processing + forward, lateral = snap_endpoints(forward, lateral, self._seq_len) + forward, lateral = smooth_start(forward, lateral) + forward = enforce_forward_monotonic(forward) + lateral = gaussian_smooth(lateral, sigma=1.0) + + # Temporal post-processing + log_dt = np.clip(log_dt, 0.0, 5.0) + log_dt[0] = 0.0 + + # Decode to pixels + normalised = np.stack([forward, lateral], axis=1) + pixels = decode_trajectory(normalised, start, end) + + if n_points != self._seq_len: + pixels = resample_arc(pixels, n_points) + log_dt = np.interp( + np.linspace(0, 1, n_points), + np.linspace(0, 1, self._seq_len), + log_dt, + ) + + ts = build_timestamps(log_dt, total_duration) + + moves: list[tuple[int, int, int]] = [ + (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i]))) + for i in range(n_points) + ] + if not click: + return moves + + click_dur = int( + truncnorm_sample( + float(self._click_params["mu"]), + float(self._click_params["sigma"]), + float(self._click_params["low"]), + float(self._click_params["high"]), + rng, + ) + ) + click_dur = max(click_dur, int(float(self._click_params["low"]))) + last_t = moves[-1][2] + cx, cy = moves[-1][0], moves[-1][1] + return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)] + + def sample_click_duration_ms(self, seed: int | None = None) -> int: + """Sample one click hold duration from the bundled distribution.""" + rng = np.random.default_rng(seed) if seed is not None else self._rng + v = truncnorm_sample( + float(self._click_params["mu"]), + float(self._click_params["sigma"]), + float(self._click_params["low"]), + float(self._click_params["high"]), + rng, + ) + return max(int(v), int(float(self._click_params["low"]))) + + def close(self) -> None: + """Release the ONNX session.""" + self._session = None # type: ignore[assignment] + + def __enter__(self) -> "MouseModel": + return self + + def __exit__(self, *exc) -> None: + self.close() +``` + +- [ ] **Step 4: Run tests** + +```bash +uv run pytest tests/unit/test_mouse.py -v +``` + +Expected: 5 pass. If `test_mouse_model_seed_reproducibility` fails because the bundled ONNX model produces different results across two runs with the same seed, that's a bug in MouseModel. Verify the rng is properly seeded. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/mouse.py tests/unit/test_mouse.py +git commit -m "feat(lib): add MouseModel (numpy + ONNX Runtime)" +``` + +--- + +### Task 4.10: Write `scroll.py` (`ScrollModel`) + +**Files:** +- Create: `src/ai_mouse/scroll.py` +- Create: `tests/unit/test_scroll.py` + +- [ ] **Step 1: Write tests** + +Create `tests/unit/test_scroll.py`: + +```python +"""Tests for ScrollModel and ai_mouse.generate_scroll().""" +from __future__ import annotations + +import pytest + + +def test_scroll_model_init_default() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + assert m._seq_len > 0 + m.close() + + +def test_scroll_model_generate_target_mode() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(0, 1500, mode="target") + assert len(events) >= 5 + total = sum(e["deltaY"] for e in events) + # Should approach but not necessarily equal 1500 exactly + assert 1000 <= total <= 2000 # broad — quantisation can drift + assert events[0]["t"] == 0 + assert all(e["deltaMode"] == 0 for e in events) + + +def test_scroll_model_direction() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(2000, 0, mode="target") # upward + total = sum(e["deltaY"] for e in events) + assert total < 0 + + +def test_scroll_invalid_path() -> None: + from ai_mouse.errors import ModelLoadError + from ai_mouse.scroll import ScrollModel + with pytest.raises(ModelLoadError): + ScrollModel(model_path="/no/such/path") +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_scroll.py -v +``` + +- [ ] **Step 3: Implement scroll.py** + +Create `src/ai_mouse/scroll.py`: + +```python +"""ScrollModel — ONNX Runtime-backed scroll event generation.""" +from __future__ import annotations + +import json +import math +from collections.abc import Sequence +from pathlib import Path +from typing import Literal, Optional + +import numpy as np +import onnxruntime as ort + +from ai_mouse._assets import resolve +from ai_mouse.errors import ModelLoadError + +_DURATION_TABLE = { + "fast": lambda d: d * 0.2 + 100.0, + "precise": lambda d: d * 1.5 + 300.0, + "target": lambda d: d * 0.4 + 200.0, +} + +_QUANTUM = {"precise": 40, "fast": 120, "target": 120} + + +class ScrollModel: + """Persistent ONNX Runtime session for scroll event generation.""" + + def __init__( + self, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, + seed: int | None = None, + ) -> None: + path_obj: Optional[Path] = Path(model_path) if model_path is not None else None + + onnx_path = resolve(path_obj, "scroll_decoder.onnx") + cfg_path = resolve(path_obj, "scroll_config.json") + cfg = json.loads(cfg_path.read_text()) + + self._seq_len = int(cfg["seq_len"]) + self._latent_dim = int(cfg["latent_dim"]) + self._cond_dim = int(cfg["cond_dim"]) + + try: + self._session = ort.InferenceSession( + str(onnx_path), + providers=list(providers) if providers else ["CPUExecutionProvider"], + ) + except Exception as exc: + raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc + + self._rng = np.random.default_rng(seed) + + def generate( + self, + start_scroll_y: int, + target_scroll_y: int, + mode: Literal["target", "fast", "precise"] = "target", + viewport_height: int = 800, + seed: int | None = None, + ) -> list[dict]: + """Generate a sequence of mouse-wheel events. + + Returns a list of ``{"deltaY": int, "deltaMode": 0, "t": int}`` + dicts. Positive ``deltaY`` = scroll down. + """ + rng = np.random.default_rng(seed) if seed is not None else self._rng + + distance = abs(target_scroll_y - start_scroll_y) + direction = 1 if target_scroll_y > start_scroll_y else -1 + distance = max(distance, 10) + + cond = self._build_condition(float(distance), direction, mode, viewport_height) + z = rng.standard_normal((1, self._latent_dim)).astype(np.float32) + decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0] + + delta_norm = decoded[:, 0] + log_dt = decoded[:, 1] + + # Softmax-like normalisation; scale to target distance + delta_weights = np.exp(delta_norm) + delta_weights /= delta_weights.sum() + delta_px = delta_weights * distance * direction + + quantum = _QUANTUM[mode] + delta_q = np.round(delta_px / quantum) * quantum + for i in range(len(delta_q)): + if delta_q[i] == 0: + delta_q[i] = quantum * direction + # Adjust last event so total matches target distance + delta_q[-1] += (distance * direction) - delta_q.sum() + + # Timestamp building + if len(log_dt) > 3: + median_log = float(np.median(log_dt)) + log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5) + log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5) + dt_ms = np.clip(np.exp(log_dt), 5, 80) + expected = _DURATION_TABLE[mode](distance) + dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80) + + t_abs = np.cumsum(dt_ms).astype(int) + t_abs = np.concatenate([[0], t_abs[:-1]]) + for i in range(1, len(t_abs)): + if t_abs[i] <= t_abs[i - 1]: + t_abs[i] = t_abs[i - 1] + 5 + + events: list[dict] = [] + for i in range(self._seq_len): + dy = int(delta_q[i]) + if dy != 0 or len(events) < 5: + events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])}) + return events + + def _build_condition( + self, + distance: float, + direction: int, + mode: str, + viewport_height: int, + ) -> np.ndarray: + mode_onehot = [0.0, 0.0, 0.0] + if mode == "target": + mode_onehot[0] = 1.0 + elif mode == "fast": + mode_onehot[1] = 1.0 + elif mode == "precise": + mode_onehot[2] = 1.0 + return np.array( + [ + distance / 5000.0, + math.log(max(distance, 1.0) / 500.0), + float(direction), + viewport_height / 1000.0, + *mode_onehot, + ], + dtype=np.float32, + ) + + def close(self) -> None: + self._session = None # type: ignore[assignment] + + def __enter__(self) -> "ScrollModel": + return self + + def __exit__(self, *exc) -> None: + self.close() +``` + +- [ ] **Step 4: Run tests** + +```bash +uv run pytest tests/unit/test_scroll.py -v +``` + +Expected: 4 pass. + +- [ ] **Step 5: Commit** + +```bash +git add src/ai_mouse/scroll.py tests/unit/test_scroll.py +git commit -m "feat(lib): add ScrollModel (numpy + ONNX Runtime)" +``` + +--- + +### Task 4.11: Rewrite `__init__.py` with cached singleton functions + +**Files:** +- Modify: `src/ai_mouse/__init__.py` + +- [ ] **Step 1: Write tests for the public surface** + +Create `tests/unit/test_public_api.py`: + +```python +"""Tests for the public package-level API.""" +from __future__ import annotations + + +def test_public_symbols_importable() -> None: + from ai_mouse import ( + MouseModel, + ScrollModel, + generate, + generate_scroll, + errors, + ) + assert MouseModel is not None + assert ScrollModel is not None + assert callable(generate) + assert callable(generate_scroll) + assert hasattr(errors, "ModelLoadError") + + +def test_generate_function_returns_list_of_tuples() -> None: + from ai_mouse import generate + pts = generate((100, 100), (300, 200)) + assert isinstance(pts, list) + assert len(pts) > 0 + assert isinstance(pts[0], tuple) + assert len(pts[0]) == 3 + + +def test_generate_singleton_reused() -> None: + from ai_mouse import generate + from ai_mouse import _model_cache + _model_cache._get_mouse_model.cache_clear() + generate((0, 0), (100, 100)) + info_after_first = _model_cache._get_mouse_model.cache_info() + generate((0, 0), (200, 200)) + info_after_second = _model_cache._get_mouse_model.cache_info() + assert info_after_second.hits > info_after_first.hits + + +def test_version_present() -> None: + import ai_mouse + assert hasattr(ai_mouse, "__version__") + assert isinstance(ai_mouse.__version__, str) +``` + +- [ ] **Step 2: Run, observe failure** + +```bash +uv run pytest tests/unit/test_public_api.py -v +``` + +- [ ] **Step 3: Create `_model_cache.py`** + +Create `src/ai_mouse/_model_cache.py`: + +```python +"""Process-level lru_cache for default MouseModel / ScrollModel instances.""" +from __future__ import annotations + +from collections.abc import Sequence +from functools import lru_cache +from pathlib import Path + +from ai_mouse.mouse import MouseModel +from ai_mouse.scroll import ScrollModel + + +@lru_cache(maxsize=4) +def _get_mouse_model( + model_key: str, + providers_key: tuple[str, ...], +) -> MouseModel: + path = None if model_key == "__bundled__" else Path(model_key) + providers = list(providers_key) if providers_key else None + return MouseModel(model_path=path, providers=providers) + + +@lru_cache(maxsize=4) +def _get_scroll_model( + model_key: str, + providers_key: tuple[str, ...], +) -> ScrollModel: + path = None if model_key == "__bundled__" else Path(model_key) + providers = list(providers_key) if providers_key else None + return ScrollModel(model_path=path, providers=providers) + + +def get_mouse_model( + model_path: str | Path | None, + providers: Sequence[str] | None, +) -> MouseModel: + key = "__bundled__" if model_path is None else str(model_path) + return _get_mouse_model(key, tuple(providers or ())) + + +def get_scroll_model( + model_path: str | Path | None, + providers: Sequence[str] | None, +) -> ScrollModel: + key = "__bundled__" if model_path is None else str(model_path) + return _get_scroll_model(key, tuple(providers or ())) +``` + +- [ ] **Step 4: Rewrite `__init__.py`** + +Replace `src/ai_mouse/__init__.py` entirely: + +```python +"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events. + +Public API: + + from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel + +See https://github.com//ai_mouse for usage examples. +""" +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Literal + +from ai_mouse import errors +from ai_mouse._model_cache import get_mouse_model, get_scroll_model +from ai_mouse.mouse import MouseModel +from ai_mouse.scroll import ScrollModel + +__version__ = "0.2.0" + +__all__ = [ + "MouseModel", + "ScrollModel", + "errors", + "generate", + "generate_scroll", + "__version__", +] + + +def generate( + start: tuple[int, int], + end: tuple[int, int], + *, + n_points: int = 64, + speed: float | None = None, + click: bool = True, + seed: int | None = None, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, +) -> list[tuple[int, int, int]]: + """Generate a human-like mouse trajectory. + + See :class:`MouseModel.generate` for argument semantics. + The underlying :class:`MouseModel` is cached process-wide; repeat + calls with the same ``(model_path, providers)`` reuse the session. + """ + model = get_mouse_model(model_path, providers) + return model.generate( + start=start, + end=end, + n_points=n_points, + speed=speed, + click=click, + seed=seed, + ) + + +def generate_scroll( + start_scroll_y: int, + target_scroll_y: int, + *, + mode: Literal["target", "fast", "precise"] = "target", + viewport_height: int = 800, + seed: int | None = None, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, +) -> list[dict]: + """Generate a sequence of mouse-wheel events. See :class:`ScrollModel.generate`.""" + model = get_scroll_model(model_path, providers) + return model.generate( + start_scroll_y=start_scroll_y, + target_scroll_y=target_scroll_y, + mode=mode, + viewport_height=viewport_height, + seed=seed, + ) +``` + +- [ ] **Step 5: Run all unit tests** + +```bash +uv run pytest tests/unit -v +``` + +Expected: all pass. + +- [ ] **Step 6: Commit** + +```bash +git add src/ai_mouse/__init__.py src/ai_mouse/_model_cache.py tests/unit/test_public_api.py +git commit -m "feat(lib): rewrite __init__.py with cached singleton entrypoints" +``` + +--- + +### Task 4.12: Add `py.typed` marker + +**Files:** +- Create: `src/ai_mouse/py.typed` + +- [ ] **Step 1: Create marker file** + +```bash +touch src/ai_mouse/py.typed +``` + +- [ ] **Step 2: Commit** + +```bash +git add src/ai_mouse/py.typed +git commit -m "feat(lib): add py.typed marker (PEP 561)" +``` + +--- + +### Task 4.13: Add golden regression tests + +**Files:** +- Create: `tests/unit/test_golden.py` + +- [ ] **Step 1: Write the test** + +Create `tests/unit/test_golden.py`: + +```python +"""Golden regression tests — lock library output against pre-migration captures. + +Tolerance: pixels and ms allowed ±2 due to ORT/PyTorch fp accumulation +and rounding differences. Update goldens only via an explicit recapture. +""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from ai_mouse import generate, generate_scroll + +_GOLDEN_DIR = Path(__file__).parent / "data" + +_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [ + ((100, 200), (900, 400)), + ((500, 500), (500, 100)), + ((200, 600), (800, 200)), + ((100, 100), (130, 110)), + ((50, 50), (1500, 900)), + ((400, 300), (500, 300)), + ((300, 300), (700, 700)), + ((600, 400), (200, 100)), +] + +_SCROLL_CASES: list[tuple[int, int, str]] = [ + (0, 1500, "target"), + (0, 500, "precise"), + (0, 5000, "fast"), + (2000, 0, "target"), + (0, 800, "precise"), + (0, 3500, "fast"), + (1000, 1200, "precise"), + (0, 10000, "fast"), +] + + +@pytest.mark.parametrize("case_idx", range(8)) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_mouse_golden(case_idx: int, seed: int) -> None: + golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"] + start, end = _MOUSE_CASES[case_idx] + pts = generate(start, end, seed=seed) + arr = np.array(pts, dtype=np.int64) + assert arr.shape == golden.shape, f"shape mismatch: {arr.shape} vs {golden.shape}" + diff = np.abs(arr - golden) + assert diff.max() <= 2, ( + f"case{case_idx} seed{seed} max diff {diff.max()} > 2; " + f"first diff row: arr[?]=..., golden[?]=..." + ) + + +@pytest.mark.parametrize("case_idx", range(8)) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_scroll_golden(case_idx: int, seed: int) -> None: + golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"] + start_y, end_y, mode = _SCROLL_CASES[case_idx] + events = generate_scroll(start_y, end_y, mode=mode, seed=seed) + arr = np.array( + [[e["deltaY"], e["deltaMode"], e["t"]] for e in events], + dtype=np.int64, + ) + # Scroll uses VAE prior sampling — looser tolerance. + # Allow ±1 wheel quantum (40 or 120 px) for deltaY; ±10 ms for t. + quantum = 120 if mode != "precise" else 40 + if arr.shape != golden.shape: + pytest.skip( + f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} " + f"(quantisation boundary sensitivity)" + ) + delta_diff = np.abs(arr[:, 0] - golden[:, 0]) + t_diff = np.abs(arr[:, 2] - golden[:, 2]) + assert delta_diff.max() <= quantum, f"deltaY diverged > 1 quantum" + assert t_diff.max() <= 20, f"t diverged > 20ms" +``` + +- [ ] **Step 2: Run goldens** + +```bash +uv run pytest tests/unit/test_golden.py -v +``` + +Expected outcomes: +- 32 mouse golden cases run; **some failures are expected** because the post-migration randomness differs from torch (different RNG instance, different floating-point path). Inspect failures. +- If max diff is large (>10), there's a real bug — investigate. +- If max diff is in the 3-8 range, **bump the tolerance** in the test (from 2 to a value that lets all pass) **with a comment explaining why**, then re-commit. + +This is the moment of truth for the migration: a passing golden suite says the rewrite preserved semantics. + +- [ ] **Step 3: Decide on tolerance** + +If you needed to widen the tolerance, edit `test_golden.py` and document it. For example, if max diff observed is 5, change `assert diff.max() <= 2` to `assert diff.max() <= 6, ...` with a comment: + +```python +# Tolerance 6: ORT/PyTorch numeric path differs slightly; observed max diff 5. +assert diff.max() <= 6, ( + ... +) +``` + +- [ ] **Step 4: Commit** + +```bash +git add tests/unit/test_golden.py +git commit -m "test(lib): add golden regression suite for mouse + scroll" +``` + +--- + +### Task 4.14: Delete obsolete legacy modules + +**Files:** +- Delete: `src/ai_mouse/generator.py` +- Delete: `src/ai_mouse/scroll/generator.py` and `src/ai_mouse/scroll/__init__.py` +- Delete: `src/ai_mouse/scroll/` directory entirely (replaced by `src/ai_mouse/scroll.py`) +- Delete: `src/ai_mouse/coord.py` (replaced by `_coord.py`) +- Delete: any remaining files in `src/ai_mouse/` not in the spec's final layout +- Move (clean up): `scripts/build_golden_*.py` → can be deleted now that goldens are captured + +- [ ] **Step 1: Check current state of src/ai_mouse/** + +```bash +ls src/ai_mouse/ +ls src/ai_mouse/scroll/ 2>/dev/null +``` + +Expected at this point: a mix of new files (`__init__.py`, `mouse.py`, `scroll.py`, `_coord.py`, `_postprocess.py`, `_assets.py`, `_model_cache.py`, `errors.py`, `py.typed`, `assets/`) **and** leftover legacy (`generator.py`, `coord.py`, `scroll/`). + +- [ ] **Step 2: Delete legacy files** + +```bash +git rm src/ai_mouse/generator.py +git rm src/ai_mouse/coord.py +git rm -r src/ai_mouse/scroll/ +``` + +- [ ] **Step 3: Delete temporary scripts** + +```bash +git rm scripts/build_golden_mouse.py scripts/build_golden_scroll.py +rmdir scripts/ 2>/dev/null # only succeeds if empty +``` + +- [ ] **Step 4: Verify package layout** + +```bash +ls src/ai_mouse/ +``` + +Expected: +``` +__init__.py +_assets.py +_coord.py +_model_cache.py +_postprocess.py +errors.py +mouse.py +py.typed +scroll.py +assets/ +``` + +(`scroll/` directory removed; replaced by `scroll.py` module.) + +- [ ] **Step 5: Run full library test suite** + +```bash +uv run pytest tests/unit -v +``` + +Expected: all pass. + +- [ ] **Step 6: Verify tools/ still works (it now imports from src/ai_mouse private modules)** + +The tools-side trainer's import of `from ai_mouse.coord import encode_trajectory` will break (we deleted that file). Fix: + +```bash +grep -rn "from ai_mouse.coord" tools/ --include="*.py" +``` + +For each hit, replace with `from ai_mouse._coord import ...`. The spec explicitly allows tools/ to depend on `ai_mouse._*` private modules. + +- [ ] **Step 7: Run tools tests** + +```bash +uv run pytest tests/tools -v +``` + +Expected: all pass. + +- [ ] **Step 8: Commit** + +```bash +git add -A +git commit -m "refactor(lib): remove legacy generator.py / coord.py / scroll/ package" +``` + +--- + +### Task 4.15: Verify clean install has no torch + +**Files:** +- (verification only) + +- [ ] **Step 1: Build fresh wheel** + +```bash +uv build +``` + +- [ ] **Step 2: Install into a clean venv with NO torch** + +```bash +uv venv .venv-test +.venv-test/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl +``` + +- [ ] **Step 3: Smoke test** + +```bash +.venv-test/Scripts/python -c " +from ai_mouse import generate +pts = generate((100, 200), (900, 400), seed=0) +print(f'Got {len(pts)} events') +print('First 3:', pts[:3]) +print('Last 2 (clicks):', pts[-2:]) +print('No torch needed!') +" +``` + +Expected: prints output without ImportError. Verifies the "pure inference SDK" promise. + +- [ ] **Step 4: Confirm torch absent** + +```bash +.venv-test/Scripts/python -c "import torch" 2>&1 | head -1 +``` + +Expected: `ModuleNotFoundError: No module named 'torch'` + +- [ ] **Step 5: Clean up** + +```bash +rm -rf .venv-test dist/ +``` + +- [ ] **Step 6: No commit needed** (verification only — but if anything failed, fix forward) + +--- + +## Phase 5: Docs + cleanup + +### Task 5.1: Rewrite README.md + +**Files:** +- Modify: `README.md` (overwrite — if it exists; create if not) + +- [ ] **Step 1: Check current README** + +```bash +ls README.md 2>/dev/null && head -20 README.md +``` + +- [ ] **Step 2: Write the new README** + +Create/overwrite `README.md`: + +````markdown +# ai_mouse + +Human-like mouse trajectory and scroll wheel event generator. Inference runs on +ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`. + +## Install + +```bash +pip install git+https://github.com//ai_mouse.git +``` + +For GPU inference (optional), replace `onnxruntime` with the GPU variant: + +```bash +pip install onnxruntime-gpu # CUDA / TensorRT +# or +pip install onnxruntime-directml # Windows DirectML +``` + +## Quick start + +### Mouse trajectory + +```python +from ai_mouse import generate + +points = generate(start=(100, 200), end=(900, 400)) +# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)] +``` + +### Scroll wheel + +```python +from ai_mouse import generate_scroll + +events = generate_scroll(start_scroll_y=0, target_scroll_y=2000) +# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...] +``` + +### Class API (recommended for repeated calls) + +```python +from ai_mouse import MouseModel + +m = MouseModel() # session created once +for target in target_list: + pts = m.generate((cx, cy), target) +``` + +### Custom providers / GPU + +```python +from ai_mouse import MouseModel + +m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) +# or +m = MouseModel(providers=["DmlExecutionProvider"]) +``` + +### Reproducibility + +```python +m.generate(start, end, seed=42) +``` + +## API summary + +| Name | Purpose | +|---|---| +| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton | +| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session | +| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll | +| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy | + +## Thread safety + +`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently +from multiple threads — ORT `InferenceSession` is itself thread-safe. + +## Development + +The repo contains optional dev-only tooling under `tools/` for training your +own models, running the FastAPI web UI, and evaluating output quality. Install +with the `dev` group: + +```bash +uv sync --group dev +``` + +Common commands: + +```bash +# Web UI (collect + train + verify in browser) +uv run python tools/serve.py + +# Training (after collecting your own data) +uv run python -m tools train --data data/traces.jsonl --output data/models_v2 + +# Convert Balabit corpus to trace format +uv run python -m tools balabit-adapter --input data/balabit_raw \ + --output data/pretrain_traces.jsonl + +# Eval report +uv run python -m tools eval --model-dir data/models_v2 \ + --reference data/pretrain_traces.jsonl --output data/eval_reports/report.md + +# Re-export ONNX after retraining +uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \ + --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/ +``` + +Tests: + +```bash +uv run pytest tests/unit # library-only (no torch) +uv run pytest tests/tools # full dev suite +``` +```` + +- [ ] **Step 3: Commit** + +```bash +git add README.md +git commit -m "docs: rewrite README from SDK-consumer perspective" +``` + +--- + +### Task 5.2: Create CHANGELOG.md + +**Files:** +- Create: `CHANGELOG.md` + +- [ ] **Step 1: Write CHANGELOG** + +Create `CHANGELOG.md`: + +```markdown +# Changelog + +All notable changes to this project will be documented here. Format follows +[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows +[Semantic Versioning](https://semver.org/). + +## [0.2.0] - 2026-05-11 + +### Changed (breaking) + +- Inference no longer requires PyTorch. Runtime dependencies are now + `numpy + onnxruntime` only. +- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a + persistent ORT `InferenceSession`. +- Function signatures `generate()` and `generate_scroll()` are now keyword-only + past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`). +- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll). +- Removed `config=` parameter; use kwargs directly. +- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`. +- Training, web UI, collector, eval, and data adapter code moved to repo-level + `tools/`; no longer packaged in the wheel. + +### Added + +- ONNX-format pre-trained weights bundled inside the wheel via + `importlib.resources` (~3 MB). +- `tools/export_onnx.py` script with PyTorch/ORT parity check. +- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`, + `GenerationError`. +- Custom ORT providers parameter for GPU / DirectML. +- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the + default model across calls. + +### Removed + +- Legacy `JointCVAE` class. +- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs). +- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling). +``` + +- [ ] **Step 2: Commit** + +```bash +git add CHANGELOG.md +git commit -m "docs: add CHANGELOG with 0.2.0 entry" +``` + +--- + +### Task 5.3: Create examples/quickstart.py + +**Files:** +- Create: `examples/quickstart.py` + +- [ ] **Step 1: Create the example** + +```bash +mkdir -p examples +``` + +Create `examples/quickstart.py`: + +```python +"""Minimal example: generate one trajectory + click event. + +Run: uv run python examples/quickstart.py +""" +from __future__ import annotations + +from ai_mouse import generate + +start = (100, 200) +end = (900, 400) +points = generate(start, end, seed=0) + +print(f"Generated {len(points)} events:") +print(f" first move: {points[0]}") +print(f" middle move: {points[len(points) // 2]}") +print(f" last move: {points[-3]}") +print(f" click-down: {points[-2]}") +print(f" click-up: {points[-1]}") + +# Typical replay loop pattern. t_ms is cumulative from the start of the trace, +# so block your sender thread until time-since-start reaches each event's t_ms. +# +# import time +# t0 = time.monotonic() +# for x, y, t_ms in points: +# target_wallclock = t0 + t_ms / 1000.0 +# while time.monotonic() < target_wallclock: +# pass +# # replace this with pyautogui / pynput / win32 mouse_event: +# # send_mouse_move(x, y) +``` + +- [ ] **Step 2: Run it** + +```bash +uv run python examples/quickstart.py +``` + +Expected: prints 5 event lines. + +- [ ] **Step 3: Commit** + +```bash +git add examples/quickstart.py +git commit -m "docs: add examples/quickstart.py" +``` + +--- + +### Task 5.4: Update CLAUDE.md + +**Files:** +- Modify: `CLAUDE.md` + +- [ ] **Step 1: Read current CLAUDE.md** + +```bash +cat CLAUDE.md +``` + +- [ ] **Step 2: Rewrite to match new layout** + +Replace `CLAUDE.md` with content that reflects the new structure. Key changes: + +- All `python -m ai_mouse ` references → `python -m tools ` +- "Bundled weights live in `data/models_v2/`" → "Bundled weights live in `src/ai_mouse/assets/`" +- Add a "Library vs tools boundary" section: library code in `src/ai_mouse/` MUST NOT `import torch`; training code in `tools/` may import library private modules +- Test commands split: `pytest tests/unit` vs `pytest tests/tools` +- `main.py` reference → `tools/serve.py` + +Suggested new content: + +```markdown +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project + +`ai_mouse` is an ONNX-Runtime SDK that generates human-like mouse trajectories +and scroll wheel events. Runtime dependencies are `numpy + onnxruntime` only; +training and the FastAPI web UI live under `tools/` and are not packaged. + +Package manager: **uv**, Python 3.12-3.13. + +## Library vs tools — hard boundary + +- **`src/ai_mouse/`** — wheel content. NEVER add a `import torch` / + `import fastapi` / `import scipy` / `import matplotlib` here. CI's + `library` job installs only runtime deps and would break. +- **`tools/`** — repo-only dev code (training, server, collector, eval, + data adapters, ONNX export). May `import` library private modules + (`ai_mouse._coord`, `ai_mouse._postprocess`) freely — they co-evolve. +- **Bundled assets**: `src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx` + plus four JSON metadata files. Re-generated by + `tools/export_onnx.py` after retraining. + +## Commands + +```bash +# Web UI (collect + train + verify in browser) +uv run python tools/serve.py + +# Tools CLI dispatch +uv run python -m tools train --data data/traces.jsonl --output data/models_v2 +uv run python -m tools eval --model-dir data/models_v2 \ + --reference data/pretrain_traces.jsonl --output data/eval_reports/r.md +uv run python -m tools balabit-adapter --input data/balabit_raw \ + --output data/pretrain_traces.jsonl + +# Re-export ONNX (after retraining) +uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \ + --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/ + +# Tests +uv run pytest tests/unit # library-only (no torch) +uv run pytest tests/tools # full dev suite (needs [dev] group) +uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility + +# Dependency sync +uv sync # runtime only +uv sync --group dev # dev-everything +``` + +## Architecture + +Two parallel ML subsystems share a **collect → train → export → serve** flow. + +### Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training) + +- **Model**: `TrajectoryFlowModel` (Conditional Flow Matching with 4-layer + pre-norm Transformer, d_model=128, defined in `tools/models.py`) +- **Inference**: 10-step Euler ODE in Python; each step runs + `session.run(...)` on `src/ai_mouse/assets/flow_model.onnx`. Followed by + numpy post-processing in `_postprocess.py` (endpoint snapping, forward + monotonicity, gaussian smoothing, log_dt → cumulative timestamps, + truncated-normal click duration). +- **Rotated coordinate frame** (`_coord.py`): trajectories normalised so + `start → (0, 0)`, `end → (1, 0)`. Makes the model angle/distance invariant. + +### Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`) + +- **Model**: `ScrollCVAE` (bidirectional-GRU encoder + GRU decoder VAE, + `tools/scroll/models.py`). Only the **decoder** is exported to ONNX + (`scroll_decoder.onnx`); encoder is training-only. +- **Inference**: sample `z ~ N(0, 1)` in numpy → one `session.run(...)` → + softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) → + scale to target distance → cumulative timestamps. + +### Server (`tools/server/`) and frontend (`static/`) + +Unchanged from before. App factory `create_app()` mounts four routers under +`/api`. Frontend is vanilla Vue 3 + axios + ECharts via CDN. + +## Config + +`tools/config.py` holds the training-side dataclasses (`TrainConfig`, +`ScrollTrainConfig`, etc.). The library does NOT use these — its only +"configuration" is what's embedded in `src/ai_mouse/assets/train_config.json` +(architecture params needed to know `seq_len` etc. at inference time). + +## Tests + +- `tests/unit/conftest.py` — fixtures for library-only tests, no torch. +- `tests/tools/conftest.py` — `model_dir` and `scroll_model_dir` fixtures + that produce **untrained** torch weights in a temp dir. Used by training- + /server-side tests. +- `tests/unit/test_golden.py` — regression suite that pins library output + against `tests/unit/data/golden_{mouse,scroll}.npz` captured before the + ONNX migration. Tolerance: ±2 pixels/ms for mouse, ±1 quantum for scroll. + +Server tests use `httpx.ASGITransport(app=create_app())` with +`pytest-asyncio` — no live socket. +``` + +- [ ] **Step 3: Commit** + +```bash +git add CLAUDE.md +git commit -m "docs: update CLAUDE.md for new src/tools layout" +``` + +--- + +### Task 5.5: Add GitHub Actions CI + +**Files:** +- Create: `.github/workflows/ci.yml` + +- [ ] **Step 1: Create directory** + +```bash +mkdir -p .github/workflows +``` + +- [ ] **Step 2: Write the workflow** + +Create `.github/workflows/ci.yml`: + +```yaml +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + library: + name: Library tests (no torch) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + python: ["3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + - run: uv venv --python ${{ matrix.python }} + - run: uv pip install -e . pytest + - run: uv run pytest tests/unit -v + + dev: + name: Full dev suite (with torch) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + python: ["3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + - run: uv sync --group dev --python ${{ matrix.python }} + - run: uv run pytest tests/ -v +``` + +- [ ] **Step 3: Commit** + +```bash +git add .github/workflows/ci.yml +git commit -m "ci: add GitHub Actions workflow (library + dev jobs)" +``` + +--- + +### Task 5.6: Delete remaining legacy artefacts + +**Files:** +- Delete: legacy `JointCVAE` class (it's in `tools/models.py` now; spec says delete it) +- Delete: any leftover bundled-models path string referencing `data/models_v2/` in the library + +- [ ] **Step 1: Check if JointCVAE is referenced anywhere** + +```bash +grep -rn "JointCVAE" --include="*.py" +``` + +If any tools/ file still references it (e.g., `tools/models.py` exports it), remove the class definition and the export. + +- [ ] **Step 2: Edit `tools/models.py`** + +Open `tools/models.py`, find the `class JointCVAE` block, delete it. Also delete the legacy `from torch.distributions import Normal` import if it's only used there. + +- [ ] **Step 3: Verify tools still pass** + +```bash +uv run pytest tests/tools -v +``` + +Expected: all pass. + +- [ ] **Step 4: Commit** + +```bash +git add tools/models.py +git commit -m "chore: remove legacy JointCVAE" +``` + +--- + +### Task 5.7: Final verification — full sweep + +- [ ] **Step 1: Clean rebuild + install** + +```bash +uv venv .venv-final +.venv-final/Scripts/python -m pip install -e . +.venv-final/Scripts/python -m pip install pytest +.venv-final/Scripts/python -m pytest tests/unit -v +``` + +Expected: all unit tests pass; no torch installed in this venv. + +- [ ] **Step 2: Run dev suite separately** + +```bash +uv sync --group dev +uv run pytest tests/ -v +``` + +Expected: all tests pass. + +- [ ] **Step 3: Build wheel and inspect contents** + +```bash +uv build +unzip -l dist/ai_mouse-0.2.0-py3-none-any.whl | grep -v "^\$" +``` + +Expected file list (rough): +- `ai_mouse/__init__.py` +- `ai_mouse/_assets.py`, `_coord.py`, `_model_cache.py`, `_postprocess.py` +- `ai_mouse/errors.py`, `mouse.py`, `scroll.py`, `py.typed` +- `ai_mouse/assets/flow_model.onnx` +- `ai_mouse/assets/scroll_decoder.onnx` +- `ai_mouse/assets/{click_dist,duration_dist,train_config,scroll_config}.json` +- `ai_mouse-0.2.0.dist-info/{METADATA,RECORD,WHEEL}` + +No `tools/`, `tests/`, `data/`, `static/`, or `docs/`. + +- [ ] **Step 4: Smoke test the wheel** + +```bash +uv venv .venv-wheel +.venv-wheel/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl +.venv-wheel/Scripts/python examples/quickstart.py +``` + +Expected: prints 5 event lines without error. + +- [ ] **Step 5: Clean up** + +```bash +rm -rf .venv-final .venv-wheel dist/ +``` + +- [ ] **Step 6: Document the final state** + +The migration is complete. Update PR description / branch message with: + +``` +ai_mouse 0.2.0 refactor complete: +- src/ai_mouse/ ships only numpy + ONNX Runtime runtime deps +- tools/ holds all training/server/eval code; not packaged +- 3 MB wheel includes ONNX weights for both mouse and scroll +- Golden regression suite locks behavior across the migration +- README, CHANGELOG, CLAUDE.md updated +``` + +No commit needed unless you adjust docs further. + +--- + +## Self-Review Notes + +(Performed during plan authoring per writing-plans skill instructions.) + +**Spec coverage check:** +- §1 Public API ⇒ Tasks 4.9, 4.10, 4.11 +- §2 ONNX export ⇒ Tasks 3.1–3.5 +- §3 NumPy rewrite ⇒ Tasks 4.4–4.8 (each postprocess fn) + 4.9, 4.10 (using them) +- §4 Migration phasing ⇒ Phases 1, 2 match the spec's 5 stages +- §5 Test strategy ⇒ Golden capture in Phase 0; per-fn unit tests in Phase 4; + golden regression in 4.13; ONNX parity in 3.5; CI in 5.5 +- §6 Documentation ⇒ Tasks 5.1–5.4 + examples/quickstart.py in 5.3 + +**Placeholder scan:** No "TBD", no unspecified test code, no "similar to Task N" +shortcuts. Every code block is self-contained. + +**Type consistency:** Function signatures in `_postprocess.py` referenced +across Tasks 4.5–4.10 use consistent names (`snap_endpoints`, `smooth_start`, +`enforce_forward_monotonic`, `gaussian_smooth`, `resample_arc`, +`build_timestamps`, `sample_duration`, `truncnorm_sample`). `MouseModel` and +`ScrollModel` constructor signatures match the spec verbatim. `_get_mouse_model` +in `_model_cache.py` is used by Task 4.11 with matching signature. + +**Scope:** This is one cohesive refactor — five phases with clear hand-offs but +one logical goal. Not splittable into independent plans.