feat(lib): switch __init__ to ONNX-backed generate/generate_scroll
This commit is contained in:
@@ -1,5 +1,78 @@
|
|||||||
# sites/ai_mouse/ai_mouse/__init__.py
|
"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events.
|
||||||
from ai_mouse.generator import generate
|
|
||||||
from ai_mouse._scroll_legacy import generate_scroll
|
|
||||||
|
|
||||||
__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,
|
||||||
|
)
|
||||||
|
|||||||
45
src/ai_mouse/_model_cache.py
Normal file
45
src/ai_mouse/_model_cache.py
Normal file
@@ -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 ()))
|
||||||
@@ -43,3 +43,45 @@ def test_mouse_model_invalid_path_raises_model_load_error() -> None:
|
|||||||
from ai_mouse.errors import ModelLoadError
|
from ai_mouse.errors import ModelLoadError
|
||||||
with pytest.raises(ModelLoadError):
|
with pytest.raises(ModelLoadError):
|
||||||
MouseModel(model_path="/nonexistent/path/here")
|
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
|
||||||
|
|||||||
43
tests/unit/test_public_api.py
Normal file
43
tests/unit/test_public_api.py
Normal file
@@ -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)
|
||||||
@@ -35,3 +35,19 @@ def test_scroll_invalid_path() -> None:
|
|||||||
from ai_mouse.scroll import ScrollModel
|
from ai_mouse.scroll import ScrollModel
|
||||||
with pytest.raises(ModelLoadError):
|
with pytest.raises(ModelLoadError):
|
||||||
ScrollModel(model_path="/no/such/path")
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user