From d39db46170fd4a263029d9bacc833dfc33c662a3 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Tue, 12 May 2026 01:16:03 +0800 Subject: [PATCH] feat(lib): switch __init__ to ONNX-backed generate/generate_scroll --- src/ai_mouse/__init__.py | 81 +++++++++++++++++++++++++++++++++-- src/ai_mouse/_model_cache.py | 45 +++++++++++++++++++ tests/unit/test_mouse.py | 42 ++++++++++++++++++ tests/unit/test_public_api.py | 43 +++++++++++++++++++ tests/unit/test_scroll.py | 16 +++++++ 5 files changed, 223 insertions(+), 4 deletions(-) create mode 100644 src/ai_mouse/_model_cache.py create mode 100644 tests/unit/test_public_api.py diff --git a/src/ai_mouse/__init__.py b/src/ai_mouse/__init__.py index 027a92a..7f2a4c5 100644 --- a/src/ai_mouse/__init__.py +++ b/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, + ) diff --git a/src/ai_mouse/_model_cache.py b/src/ai_mouse/_model_cache.py new file mode 100644 index 0000000..dd88e7e --- /dev/null +++ b/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 ())) diff --git a/tests/unit/test_mouse.py b/tests/unit/test_mouse.py index 0488af7..ff8a5af 100644 --- a/tests/unit/test_mouse.py +++ b/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 diff --git a/tests/unit/test_public_api.py b/tests/unit/test_public_api.py new file mode 100644 index 0000000..a6a7f19 --- /dev/null +++ b/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) diff --git a/tests/unit/test_scroll.py b/tests/unit/test_scroll.py index 93eb36c..41ed84b 100644 --- a/tests/unit/test_scroll.py +++ b/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)