From 3c1cf171a79685e9bebe5957cdbbb4ab3e18951e Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:04:37 +0800
Subject: [PATCH] feat(lib): add _assets module for bundled-weight resolution

Provide bundled_path() and resolve() helpers that locate ONNX
weights and JSON metadata via importlib.resources, falling back to
a user-supplied directory. Missing assets raise ModelLoadError.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 src/ai_mouse/_assets.py   | 52 +++++++++++++++++++++++++++++++++++++++
 tests/unit/test_assets.py | 34 +++++++++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 src/ai_mouse/_assets.py
 create mode 100644 tests/unit/test_assets.py

diff --git a/src/ai_mouse/_assets.py b/src/ai_mouse/_assets.py
new file mode 100644
index 0000000..cf5c6f5
--- /dev/null
+++ b/src/ai_mouse/_assets.py
@@ -0,0 +1,52 @@
+"""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`. For non-zip wheels (the common
+        case) the path points directly into the installed package; for
+        zipapp installations :func:`importlib.resources.as_file`
+        materialises a temp file.
+    """
+    ref = files(_PACKAGE_ASSETS) / name
+    with as_file(ref) as p:
+        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
diff --git a/tests/unit/test_assets.py b/tests/unit/test_assets.py
new file mode 100644
index 0000000..c52512a
--- /dev/null
+++ b/tests/unit/test_assets.py
@@ -0,0 +1,34 @@
+"""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")