geointel/backend/app/services/model_asset_catalog_service.py

from __future__ import annotations

import hashlib
import re
from pathlib import Path

from app.core.config import Settings, get_settings
from app.core.errors import AppError
from app.schemas.detection import ModelAssetListResponse, ModelAssetRead


class ModelAssetCatalogService:
    SUPPORTED_SUFFIXES = {
        ".pt": "ultralytics/pytorch",
        ".onnx": "onnx",
        ".engine": "tensorrt",
    }

    @staticmethod
    def list_assets(settings: Settings | None = None) -> ModelAssetListResponse:
        resolved_settings = settings or get_settings()
        model_directory = ModelAssetCatalogService._model_directory(resolved_settings)
        active_model_path = ModelAssetCatalogService._resolved_file_path(resolved_settings.yolo_model_path)
        if not model_directory.exists() or not model_directory.is_dir():
            return ModelAssetListResponse(items=[], total=0, model_directory=str(model_directory))

        candidate_paths = [
            path
            for path in sorted(model_directory.iterdir(), key=lambda item: item.name.lower())
            if path.is_file() and path.suffix.lower() in ModelAssetCatalogService.SUPPORTED_SUFFIXES
        ]
        if active_model_path is not None:
            candidate_paths = [path for path in candidate_paths if path.resolve() == active_model_path]

        items = [
            ModelAssetCatalogService._asset_from_file(path, active_model_path=active_model_path)
            for path in candidate_paths
        ]
        return ModelAssetListResponse(items=items, total=len(items), model_directory=str(model_directory))

    @staticmethod
    def resolve_asset(model_asset_id: str, settings: Settings | None = None) -> ModelAssetRead:
        normalized = model_asset_id.strip()
        for asset in ModelAssetCatalogService.list_assets(settings=settings).items:
            if asset.model_asset_id == normalized:
                return asset
        raise AppError(
            code="DETECTION_MODEL_ASSET_NOT_FOUND",
            message="Selected local model asset was not found in the configured model directory",
            details={"model_asset_id": normalized},
            status_code=404,
        )

    @staticmethod
    def settings_for_asset(settings: Settings, asset: ModelAssetRead) -> Settings:
        return settings.model_copy(update={"yolo_model_path": asset.model_path})

    @staticmethod
    def _model_directory(settings: Settings) -> Path:
        configured_directory = Path(settings.yolo_models_dir).expanduser()
        if configured_directory.exists() and configured_directory.is_dir():
            return configured_directory.resolve()
        active_model_path = ModelAssetCatalogService._resolved_file_path(settings.yolo_model_path)
        if active_model_path and active_model_path.parent.exists() and active_model_path.parent.is_dir():
            return active_model_path.parent.resolve()
        return configured_directory.resolve()

    @staticmethod
    def _asset_from_file(path: Path, *, active_model_path: Path | None) -> ModelAssetRead:
        resolved_path = path.resolve()
        return ModelAssetRead(
            model_asset_id=ModelAssetCatalogService._asset_id(path),
            filename=path.name,
            display_name=path.stem,
            model_path=str(resolved_path),
            suffix=path.suffix.lower(),
            framework=ModelAssetCatalogService.SUPPORTED_SUFFIXES[path.suffix.lower()],
            task_type="object_detection",
            size_bytes=path.stat().st_size,
            sha256=ModelAssetCatalogService._sha256(path),
            active=active_model_path == resolved_path,
            runtime_available=True,
            runtime_status="active" if active_model_path == resolved_path else "available",
            governed_validation_status="not_verified_by_catalog",
            promotion_status="not_verified_by_catalog",
            status="runtime_active" if active_model_path == resolved_path else "runtime_available",
            limitation_message=(
                "Active local runtime model asset. Runtime selection is not evidence of governed validation or promotion."
                if active_model_path == resolved_path
                else "Local runtime model asset. Governed validation and promotion are not established by this catalog."
            ),
            will_download_models=False,
        )

    @staticmethod
    def _asset_id(path: Path) -> str:
        raw = f"{path.stem}-{path.suffix.lower().lstrip('.')}"
        normalized = re.sub(r"[^a-z0-9]+", "-", raw.lower()).strip("-")
        return normalized or "model-asset"

    @staticmethod
    def _sha256(path: Path) -> str:
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        return digest.hexdigest()

    @staticmethod
    def _resolved_file_path(raw_path: str | None) -> Path | None:
        if not raw_path:
            return None
        path = Path(raw_path).expanduser()
        if not path.exists() or not path.is_file():
            return None
        return path.resolve()