feat(lib): switch __init__ to ONNX-backed generate/generate_scroll

src/ai_mouse/__init__.py

@@ -1,5 +1,78 @@
-# sites/ai_mouse/ai_mouse/__init__.py
-from ai_mouse.generator import generate
-from ai_mouse._scroll_legacy import generate_scroll
+"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events.
 
-__all__ = ["generate", "generate_scroll"]
+Public API:
+
+    from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel
+
+See the project README 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 :meth:`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,
+    )

src/ai_mouse/_model_cache.py

@@ -0,0 +1,45 @@
+"""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 ()))

tests/unit/test_mouse.py

@@ -43,3 +43,45 @@ def test_mouse_model_invalid_path_raises_model_load_error() -> None:
     from ai_mouse.errors import ModelLoadError
     with pytest.raises(ModelLoadError):
         MouseModel(model_path="/nonexistent/path/here")
+
+
+def test_mouse_model_timestamps_monotonic() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    times = [p[2] for p in pts]
+    for i in range(1, len(times)):
+        assert times[i] >= times[i - 1]
+
+
+def test_mouse_model_starts_and_ends_near_endpoints() -> None:
+    from ai_mouse.mouse import MouseModel
+    start = (100, 200)
+    end = (900, 400)
+    m = MouseModel()
+    pts = m.generate(start, end)
+    assert abs(pts[0][0] - start[0]) < 30
+    assert abs(pts[0][1] - start[1]) < 30
+    last_move = pts[-3]  # last 2 are click events
+    assert abs(last_move[0] - end[0]) < 30
+    assert abs(last_move[1] - end[1]) < 30
+
+
+def test_mouse_model_n_points_parameter() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400), n_points=32)
+    # 32 moves + 2 clicks
+    assert len(pts) == 34
+
+
+def test_mouse_model_click_events_have_matching_coords() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    down, up = pts[-2], pts[-1]
+    assert down[0] == up[0]
+    assert down[1] == up[1]
+    assert up[2] > down[2]
+    # Within click_dist bounds 20..500
+    assert 20 <= up[2] - down[2] <= 500

tests/unit/test_public_api.py

@@ -0,0 +1,43 @@
+"""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)

tests/unit/test_scroll.py

@@ -35,3 +35,19 @@ def test_scroll_invalid_path() -> None:
     from ai_mouse.scroll import ScrollModel
     with pytest.raises(ModelLoadError):
         ScrollModel(model_path="/no/such/path")
+
+
+def test_scroll_model_timestamps_monotonic() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(0, 2000, mode="target")
+    times = [e["t"] for e in events]
+    for i in range(1, len(times)):
+        assert times[i] >= times[i - 1]
+
+
+def test_scroll_model_deltaY_are_integers() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(0, 2000, mode="target")
+    assert all(isinstance(e["deltaY"], int) for e in events)