nat-engine 1__py3-none-any.whl

mannf/product/uninstall.py

@@ -0,0 +1,169 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Uninstall command for NAT — removes config, data, and shell completions.
+
+Entry point::
+
+    def run_uninstall(args: argparse.Namespace) -> int:
+        ...
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Exit codes
+# ---------------------------------------------------------------------------
+_EXIT_OK = 0
+_EXIT_ABORTED = 1
+_EXIT_ERROR = 2
+
+# ---------------------------------------------------------------------------
+# Shell completion file locations
+# ---------------------------------------------------------------------------
+
+_BASH_COMPLETION_PATHS = [
+    Path("/etc/bash_completion.d/nat"),
+    Path("/usr/local/etc/bash_completion.d/nat"),
+]
+
+_ZSH_COMPLETION_PATHS = [
+    Path("/usr/local/share/zsh/site-functions/_nat"),
+    Path("/usr/share/zsh/site-functions/_nat"),
+]
+
+_FISH_COMPLETION_PATHS = [
+    Path.home() / ".config" / "fish" / "completions" / "nat.fish",
+    Path("/usr/share/fish/completions/nat.fish"),
+]
+
+
+def _remove_path(path: Path, removed: list[str], errors: list[str]) -> None:
+    """Remove *path* (file or directory) and record the outcome."""
+    try:
+        if path.is_dir():
+            shutil.rmtree(path)
+            removed.append(str(path))
+        elif path.exists():
+            path.unlink()
+            removed.append(str(path))
+    except OSError as exc:
+        errors.append(f"{path}: {exc}")
+
+
+def run_uninstall(args: argparse.Namespace) -> int:
+    """Remove NAT configuration files, data directory, and shell completions.
+
+    Parameters
+    ----------
+    args:
+        Parsed CLI arguments.  The following attributes are read:
+
+        - ``keep_config`` (bool) — preserve ``.natrc`` files.
+        - ``yes`` (bool) — skip confirmation prompt.
+        - ``purge`` (bool) — also run ``alembic downgrade base`` and remove
+          Docker volumes when ``DATABASE_URL`` is set.
+
+    Returns
+    -------
+    int
+        Exit code (0 = success, non-zero = error).
+    """
+    keep_config: bool = getattr(args, "keep_config", False)
+    skip_confirm: bool = getattr(args, "yes", False)
+    purge: bool = getattr(args, "purge", False)
+
+    print()
+    print("  ⚠️  NAT Uninstall")
+    print("  " + "─" * 48)
+
+    # -----------------------------------------------------------------------
+    # Confirmation
+    # -----------------------------------------------------------------------
+    if not skip_confirm:
+        try:
+            answer = input("  Are you sure you want to uninstall NAT? [y/N] ").strip().lower()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            print("  Aborted.")
+            return _EXIT_ABORTED
+        if answer not in ("y", "yes"):
+            print("  Aborted.")
+            return _EXIT_ABORTED
+
+    removed: list[str] = []
+    errors: list[str] = []
+
+    # -----------------------------------------------------------------------
+    # Remove .natrc files
+    # -----------------------------------------------------------------------
+    if not keep_config:
+        for natrc in (Path.home() / ".natrc", Path(".natrc")):
+            _remove_path(natrc, removed, errors)
+    else:
+        print("  ℹ️  Skipping .natrc removal (--keep-config).")
+
+    # -----------------------------------------------------------------------
+    # Remove ~/.nat/ data directory
+    # -----------------------------------------------------------------------
+    nat_data_dir = Path.home() / ".nat"
+    if nat_data_dir.exists():
+        _remove_path(nat_data_dir, removed, errors)
+
+    # -----------------------------------------------------------------------
+    # Remove shell completion files
+    # -----------------------------------------------------------------------
+    for path in (*_BASH_COMPLETION_PATHS, *_ZSH_COMPLETION_PATHS, *_FISH_COMPLETION_PATHS):
+        if path.exists():
+            _remove_path(path, removed, errors)
+
+    # -----------------------------------------------------------------------
+    # Purge: alembic downgrade base
+    # -----------------------------------------------------------------------
+    if purge:
+        db_url = os.environ.get("DATABASE_URL")
+        if db_url:
+            print("\n  🗄️  Running alembic downgrade base …")
+            result = subprocess.run(  # noqa: S603
+                ["alembic", "downgrade", "base"],
+                check=False,
+            )
+            if result.returncode != 0:
+                print("  ❌  alembic downgrade failed.", file=sys.stderr)
+                errors.append("alembic downgrade base")
+            else:
+                print("  ✅  Database tables removed.")
+        else:
+            print("  ℹ️  DATABASE_URL not set — skipping alembic downgrade.")
+
+    # -----------------------------------------------------------------------
+    # Summary
+    # -----------------------------------------------------------------------
+    print()
+    if removed:
+        print("  ✅  Removed:")
+        for item in removed:
+            print(f"       {item}")
+    else:
+        print("  ℹ️  Nothing to remove.")
+
+    if errors:
+        print("\n  ⚠️  Errors encountered:")
+        for err in errors:
+            print(f"       {err}", file=sys.stderr)
+
+    print()
+    print("  📦  To complete uninstallation, run:")
+    print("       pip uninstall nat-engine")
+    print()
+
+    return _EXIT_OK if not errors else _EXIT_ERROR

mannf/product/upgrade.py

@@ -0,0 +1,139 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Self-upgrade command for NAT — checks PyPI and upgrades the installed package.
+
+Entry point::
+
+    def run_upgrade(args: argparse.Namespace) -> int:
+        ...
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+
+_PYPI_JSON_URL = "https://pypi.org/pypi/nat-engine/json"
+
+# ---------------------------------------------------------------------------
+# Exit codes (reuse CLI constants pattern)
+# ---------------------------------------------------------------------------
+_EXIT_OK = 0
+_EXIT_ERROR = 1
+
+
+def _get_latest_pypi_version(allow_pre: bool = False) -> str | None:
+    """Query PyPI JSON API and return the latest version string.
+
+    Parameters
+    ----------
+    allow_pre:
+        When ``True``, consider pre-release versions as well.
+
+    Returns
+    -------
+    str | None
+        Latest version string, or ``None`` if the query fails.
+    """
+    try:
+        import httpx  # noqa: PLC0415
+
+        response = httpx.get(_PYPI_JSON_URL, timeout=10)
+        response.raise_for_status()
+        data = response.json()
+        if allow_pre:
+            # Return the overall latest release (may be pre-release)
+            versions = list(data.get("releases", {}).keys())
+            if not versions:
+                return None
+            # Sort by packaging.version if available, otherwise lexicographic
+            try:
+                from packaging.version import Version  # noqa: PLC0415
+
+                versions.sort(key=Version)
+            except Exception:
+                versions.sort()
+            return versions[-1]
+        else:
+            return data.get("info", {}).get("version")
+    except Exception as exc:  # noqa: BLE001
+        print(f"  ⚠️  Could not query PyPI: {exc}", file=sys.stderr)
+        return None
+
+
+def run_upgrade(args: argparse.Namespace) -> int:
+    """Check for updates and optionally upgrade ``nat-engine``.
+
+    Parameters
+    ----------
+    args:
+        Parsed CLI arguments.  The following attributes are read:
+
+        - ``check`` (bool) — only check for an update, do not install.
+        - ``migrate`` (bool) — run ``alembic upgrade head`` after upgrading.
+        - ``pre`` (bool) — allow pre-release versions when checking / installing.
+
+    Returns
+    -------
+    int
+        Exit code (0 = success, non-zero = error).
+    """
+    from mannf._version import __version__ as current_version  # noqa: PLC0415
+
+    check_only: bool = getattr(args, "check", False)
+    run_migrate: bool = getattr(args, "migrate", False)
+    allow_pre: bool = getattr(args, "pre", False)
+
+    print()
+    print(f"  🔍  Current version : nat-engine {current_version}")
+
+    latest = _get_latest_pypi_version(allow_pre=allow_pre)
+    if latest is None:
+        print("  ❌  Unable to determine the latest version from PyPI.", file=sys.stderr)
+        return _EXIT_ERROR
+
+    print(f"  🌐  Latest version  : nat-engine {latest}")
+
+    if current_version == latest:
+        print("  ✅  You are already on the latest version.")
+        print()
+        return _EXIT_OK
+
+    if check_only:
+        print(f"  ℹ️  An update is available: {current_version} → {latest}")
+        print("  Run `nat upgrade` (without --check) to install it.")
+        print()
+        return _EXIT_OK
+
+    # -----------------------------------------------------------------------
+    # Perform the upgrade
+    # -----------------------------------------------------------------------
+    pip_cmd = [sys.executable, "-m", "pip", "install", "--upgrade", "nat-engine"]
+    if allow_pre:
+        pip_cmd.append("--pre")
+
+    print(f"\n  ⬆️  Upgrading nat-engine {current_version} → {latest} …")
+    result = subprocess.run(pip_cmd, check=False)  # noqa: S603
+    if result.returncode != 0:
+        print("  ❌  pip upgrade failed.", file=sys.stderr)
+        return _EXIT_ERROR
+
+    # -----------------------------------------------------------------------
+    # Optional DB migration
+    # -----------------------------------------------------------------------
+    if run_migrate:
+        print("\n  🗄️  Running database migrations (alembic upgrade head) …")
+        alembic_cmd = ["alembic", "upgrade", "head"]
+        migrate_result = subprocess.run(alembic_cmd, check=False)  # noqa: S603
+        if migrate_result.returncode != 0:
+            print("  ❌  alembic migration failed.", file=sys.stderr)
+            return _EXIT_ERROR
+        print("  ✅  Database migrations complete.")
+
+    print(f"\n  🎉  nat-engine upgraded: {current_version} → {latest}")
+    print()
+    return _EXIT_OK

mannf/product/weights/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Weight persistence and transfer learning package."""
+
+from mannf.product.weights.store import WeightStore
+from mannf.product.weights.registry import WeightRegistry
+from mannf.product.weights.blob_store import BlobWeightStore
+from mannf.product.weights.factory import get_weight_store
+
+__all__ = ["WeightStore", "WeightRegistry", "BlobWeightStore", "get_weight_store"]

mannf/product/weights/blob_store.py

@@ -0,0 +1,299 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Azure Blob Storage backend for weight persistence.
+
+Activated when ``AZURE_STORAGE_CONNECTION_STRING`` or
+``AZURE_STORAGE_ACCOUNT_NAME`` is set in the environment.  Falls back to
+the local :class:`~mannf.product.weights.store.WeightStore` when neither
+variable is present.
+
+The blob JSON format is identical to the on-disk format used by
+:class:`~mannf.product.weights.store.WeightStore` so the two backends are
+fully interchangeable.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from datetime import datetime, timezone
+from typing import Any, Dict, List
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+_FORMAT_VERSION = "1.0"
+
+
+def _layers_to_serialisable(layers: List[dict]) -> List[dict]:
+    """Convert ``{"W": ndarray, "b": ndarray}`` layer dicts to JSON-safe lists."""
+    return [{"W": layer["W"].tolist(), "b": layer["b"].tolist()} for layer in layers]
+
+
+def _layers_from_serialisable(layers: List[dict]) -> List[dict]:
+    """Restore layer dicts from JSON-loaded lists back to ``np.ndarray`` values."""
+    return [
+        {"W": np.array(layer["W"], dtype=float), "b": np.array(layer["b"], dtype=float)}
+        for layer in layers
+    ]
+
+
+class BlobWeightStore:
+    """Azure Blob Storage backend for weight snapshots.
+
+    Activated when ``AZURE_STORAGE_CONNECTION_STRING`` or
+    ``AZURE_STORAGE_ACCOUNT_NAME`` is set.  Falls back to the local
+    :class:`~mannf.product.weights.store.WeightStore` when not configured.
+
+    Parameters
+    ----------
+    container_name:
+        Name of the Azure Blob Storage container to use.  Defaults to
+        ``"nat-weights"``.  The container is created automatically on first
+        use if it does not already exist.
+    """
+
+    def __init__(self, container_name: str = "nat-weights") -> None:
+        self._container_name = container_name
+        self._client = self._build_client()
+
+    # ------------------------------------------------------------------
+    # Public async API
+    # ------------------------------------------------------------------
+
+    async def save(self, blob_path: str, agent_weights: Dict[str, Any]) -> str:
+        """Upload weight JSON to blob storage.
+
+        Parameters
+        ----------
+        blob_path:
+            The blob name / path inside the container (e.g.
+            ``"tenant-a/my-snapshot.json"``).
+        agent_weights:
+            Mapping of ``agent_id`` → ``{"role": str, "layers": List[dict]}``.
+            Each layer dict must have ``"W"`` and ``"b"`` numpy arrays.
+
+        Returns
+        -------
+        str
+            The URL of the uploaded blob.
+        """
+        agents_serialisable: Dict[str, Any] = {}
+        for agent_id, info in agent_weights.items():
+            agents_serialisable[agent_id] = {
+                "role": info.get("role", ""),
+                "layers": _layers_to_serialisable(info["layers"]),
+            }
+
+        payload = {
+            "version": _FORMAT_VERSION,
+            "saved_at": datetime.now(timezone.utc).isoformat(),
+            "agents": agents_serialisable,
+        }
+        data = json.dumps(payload, indent=2).encode("utf-8")
+
+        blob_client = self._client.get_blob_client(
+            container=self._container_name, blob=blob_path
+        )
+        await self._ensure_container()
+        await blob_client.upload_blob(data, overwrite=True)
+        url = blob_client.url
+        logger.info("BlobWeightStore: uploaded %s (%d agents)", blob_path, len(agent_weights))
+        return url
+
+    async def load(self, blob_path: str) -> Dict[str, Any]:
+        """Download and parse weight JSON from blob storage.
+
+        Parameters
+        ----------
+        blob_path:
+            The blob name / path inside the container.
+
+        Returns
+        -------
+        dict
+            ``{"version": str, "saved_at": str, "agents": {agent_id: {...}}}``
+            where each layer dict contains ``np.ndarray`` values.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the blob does not exist.
+        ValueError
+            If the blob content cannot be parsed or fails validation.
+        """
+        from azure.core.exceptions import ResourceNotFoundError
+
+        blob_client = self._client.get_blob_client(
+            container=self._container_name, blob=blob_path
+        )
+        try:
+            stream = await blob_client.download_blob()
+            raw = await stream.readall()
+        except ResourceNotFoundError as exc:
+            raise FileNotFoundError(f"Weight blob not found: {blob_path!r}") from exc
+
+        try:
+            data = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Corrupt weight blob {blob_path!r}: {exc}") from exc
+
+        _validate_structure(data)
+
+        agents_out: Dict[str, Any] = {}
+        for agent_id, info in data["agents"].items():
+            agents_out[agent_id] = {
+                "role": info.get("role", ""),
+                "layers": _layers_from_serialisable(info["layers"]),
+            }
+
+        return {
+            "version": data["version"],
+            "saved_at": data["saved_at"],
+            "agents": agents_out,
+        }
+
+    async def delete(self, blob_path: str) -> None:
+        """Delete a weight blob.
+
+        Parameters
+        ----------
+        blob_path:
+            The blob name / path inside the container.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the blob does not exist.
+        """
+        from azure.core.exceptions import ResourceNotFoundError
+
+        blob_client = self._client.get_blob_client(
+            container=self._container_name, blob=blob_path
+        )
+        try:
+            await blob_client.delete_blob()
+        except ResourceNotFoundError as exc:
+            raise FileNotFoundError(f"Weight blob not found: {blob_path!r}") from exc
+        logger.info("BlobWeightStore: deleted %s", blob_path)
+
+    async def list_blobs(self, prefix: str = "") -> List[str]:
+        """List weight blobs with an optional prefix filter.
+
+        Parameters
+        ----------
+        prefix:
+            If non-empty, only blobs whose names begin with *prefix* are
+            returned.
+
+        Returns
+        -------
+        list[str]
+            Blob names (paths) within the container.
+        """
+        container_client = self._client.get_container_client(self._container_name)
+        blobs: List[str] = []
+        try:
+            async for blob in container_client.list_blobs(name_starts_with=prefix or None):
+                blobs.append(blob.name)
+        except Exception:  # container may not exist yet
+            pass
+        return blobs
+
+    # ------------------------------------------------------------------
+    # Class-level helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def is_configured() -> bool:
+        """Return ``True`` when Azure Blob Storage environment variables are set."""
+        return bool(
+            os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
+            or os.environ.get("AZURE_STORAGE_ACCOUNT_NAME")
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _build_client(self):  # type: ignore[return]
+        """Construct an ``AsyncBlobServiceClient`` from environment variables."""
+        try:
+            from azure.storage.blob.aio import BlobServiceClient
+        except ImportError as exc:
+            raise ImportError(
+                "azure-storage-blob is required for BlobWeightStore. "
+                "Install it with: pip install 'mannf[azure-storage]'"
+            ) from exc
+
+        conn_str = os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
+        if conn_str:
+            return BlobServiceClient.from_connection_string(conn_str)
+
+        account_name = os.environ.get("AZURE_STORAGE_ACCOUNT_NAME")
+        if account_name:
+            account_key = os.environ.get("AZURE_STORAGE_ACCOUNT_KEY", "")
+            if account_key:
+                conn_str = (
+                    f"DefaultEndpointsProtocol=https;"
+                    f"AccountName={account_name};"
+                    f"AccountKey={account_key};"
+                    f"EndpointSuffix=core.windows.net"
+                )
+                return BlobServiceClient.from_connection_string(conn_str)
+            # Managed identity / passwordless
+            from azure.identity.aio import DefaultAzureCredential
+            credential = DefaultAzureCredential()
+            return BlobServiceClient(
+                account_url=f"https://{account_name}.blob.core.windows.net",
+                credential=credential,
+            )
+
+        raise EnvironmentError(
+            "Neither AZURE_STORAGE_CONNECTION_STRING nor "
+            "AZURE_STORAGE_ACCOUNT_NAME is set."
+        )
+
+    async def _ensure_container(self) -> None:
+        """Create the container if it does not already exist."""
+        from azure.core.exceptions import ResourceExistsError
+
+        container_client = self._client.get_container_client(self._container_name)
+        try:
+            await container_client.create_container()
+            logger.debug("BlobWeightStore: created container %r", self._container_name)
+        except ResourceExistsError:
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Shared validation helper (mirrors WeightStore._validate_structure)
+# ---------------------------------------------------------------------------
+
+
+def _validate_structure(data: Any) -> None:
+    """Raise :class:`ValueError` if *data* does not match the expected schema."""
+    if not isinstance(data, dict):
+        raise ValueError("Weight blob root must be a JSON object")
+    if data.get("version") != _FORMAT_VERSION:
+        raise ValueError(
+            f"Unsupported weight blob version {data.get('version')!r} "
+            f"(expected {_FORMAT_VERSION!r})"
+        )
+    if "agents" not in data or not isinstance(data["agents"], dict):
+        raise ValueError("Weight blob missing 'agents' mapping")
+    for agent_id, info in data["agents"].items():
+        if not isinstance(info, dict):
+            raise ValueError(f"Agent entry {agent_id!r} is not a dict")
+        if "layers" not in info or not isinstance(info["layers"], list):
+            raise ValueError(f"Agent {agent_id!r} missing 'layers' list")
+        for i, layer in enumerate(info["layers"]):
+            if not isinstance(layer, dict) or "W" not in layer or "b" not in layer:
+                raise ValueError(
+                    f"Agent {agent_id!r} layer {i} must have 'W' and 'b' keys"
+                )

mannf/product/weights/factory.py

@@ -0,0 +1,42 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Factory for selecting the appropriate weight-store backend.
+
+Call :func:`get_weight_store` to obtain a store instance.  The function
+returns a :class:`~mannf.product.weights.blob_store.BlobWeightStore` when
+Azure Blob Storage environment variables are present, or the local
+:class:`~mannf.product.weights.store.WeightStore` otherwise.
+"""
+
+from __future__ import annotations
+
+from typing import Union
+
+from mannf.product.weights.blob_store import BlobWeightStore
+from mannf.product.weights.store import WeightStore
+
+
+def get_weight_store(container_name: str = "nat-weights") -> Union[BlobWeightStore, WeightStore]:
+    """Return the appropriate weight-store backend.
+
+    If ``AZURE_STORAGE_CONNECTION_STRING`` or ``AZURE_STORAGE_ACCOUNT_NAME``
+    is set in the environment, a :class:`BlobWeightStore` is returned.
+    Otherwise the local :class:`WeightStore` is returned (default behaviour,
+    fully preserves existing local-file semantics).
+
+    Parameters
+    ----------
+    container_name:
+        Passed through to :class:`BlobWeightStore` when Azure is configured.
+        Ignored for the local backend.
+
+    Returns
+    -------
+    BlobWeightStore | WeightStore
+    """
+    if BlobWeightStore.is_configured():
+        return BlobWeightStore(container_name=container_name)
+    return WeightStore()