diff --git a/python/src/coreai_models/export/bundle.py b/python/src/coreai_models/export/bundle.py index 94241ca..596f1d0 100644 --- a/python/src/coreai_models/export/bundle.py +++ b/python/src/coreai_models/export/bundle.py @@ -5,6 +5,7 @@ """Create model bundles from exported .aimodel files.""" +import importlib.metadata import json import logging from datetime import datetime @@ -17,6 +18,26 @@ METADATA_VERSION = "0.2" +_BUILD_INFO_PACKAGES = [ + "coreai-core", + "coreai-torch", + "coreai-opt", + "coreai-models", + "torch", + "transformers", +] + + +def _get_build_info() -> dict[str, str]: + """Collect versions of key packages used during export.""" + versions: dict[str, str] = {} + for pkg in _BUILD_INFO_PACKAGES: + try: + versions[pkg] = importlib.metadata.version(pkg) + except importlib.metadata.PackageNotFoundError: + pass + return versions + def bundle_llm_asset( bundle_path: Path, @@ -67,6 +88,7 @@ def _write_metadata( "date": datetime.now().astimezone().isoformat(), "targets": [], }, + "build_info": _get_build_info(), } metadata_path = bundle_path / "metadata.json" with open(metadata_path, "w") as f: diff --git a/python/tests/test_model_units/test_export/test_bundle.py b/python/tests/test_model_units/test_export/test_bundle.py new file mode 100644 index 0000000..df99eb8 --- /dev/null +++ b/python/tests/test_model_units/test_export/test_bundle.py @@ -0,0 +1,16 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +from coreai_models.export.bundle import _get_build_info + + +def test_build_info_includes_core_packages(): + info = _get_build_info() + assert "coreai-core" in info + assert "coreai-torch" in info + assert "coreai-opt" in info + assert "torch" in info + for version in info.values(): + assert version # non-empty string