yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
ai_mouse/docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md
Huang Qi 8003ae657a 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 <noreply@anthropic.com>
2026-05-12 00:01:37 +08:00

4003 lines
109 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 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: <FastAPI ...>` 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/<owner>/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:
<output>/flow_model.onnx
<output>/scroll_decoder.onnx
<output>/click_dist.json
<output>/duration_dist.json
<output>/train_config.json
<output>/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/<owner>/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/<owner>/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 <cmd>` references → `python -m tools <cmd>`
- "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.
Reference in New Issue View Git Blame Copy Permalink