diff --git a/PROJECTS/advanced/ai-threat-detection/backend/ml/metadata.py b/PROJECTS/advanced/ai-threat-detection/backend/ml/metadata.py new file mode 100644 index 00000000..f7a23bed --- /dev/null +++ b/PROJECTS/advanced/ai-threat-detection/backend/ml/metadata.py @@ -0,0 +1,82 @@ +""" +©AngelaMos | 2026 +metadata.py +""" + +import hashlib +import logging +from pathlib import Path + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.models.model_metadata import ModelMetadata + +logger = logging.getLogger(__name__) + +MODEL_TYPES: dict[str, str] = { + "ae.onnx": "autoencoder", + "rf.onnx": "random_forest", + "if.onnx": "isolation_forest", +} + +VERSION_HASH_LENGTH = 12 + + +def compute_model_version(artifact_path: Path) -> str: + """ + Compute a 12-char hex version string from the SHA-256 hash of a file + """ + sha = hashlib.sha256(artifact_path.read_bytes()) + return sha.hexdigest()[:VERSION_HASH_LENGTH] + + +async def save_model_metadata( + session: AsyncSession, + model_dir: Path, + training_samples: int, + metrics: dict[str, object], + mlflow_run_id: str | None = None, + threshold: float | None = None, +) -> list[ModelMetadata]: + """ + Persist metadata for all 3 model types, deactivating previous active versions + """ + rows: list[ModelMetadata] = [] + + for filename, model_type in MODEL_TYPES.items(): + artifact_path = model_dir / filename + version = compute_model_version(artifact_path) + + result = await session.execute( + select(ModelMetadata).where( + ModelMetadata.model_type == model_type, + ModelMetadata.is_active == True, # noqa: E712 + ) + ) + for old in result.scalars().all(): + await session.delete(old) + await session.flush() + + row = ModelMetadata( + model_type=model_type, + version=version, + training_samples=training_samples, + metrics=metrics, + artifact_path=str(artifact_path), + is_active=True, + mlflow_run_id=mlflow_run_id, + threshold=threshold, + ) + session.add(row) + rows.append(row) + + await session.commit() + + logger.info( + "Saved metadata for %d models (samples=%d)", + len(rows), + training_samples, + ) + + return rows diff --git a/PROJECTS/advanced/ai-threat-detection/backend/tests/test_metadata.py b/PROJECTS/advanced/ai-threat-detection/backend/tests/test_metadata.py new file mode 100644 index 00000000..cfb7b7ed --- /dev/null +++ b/PROJECTS/advanced/ai-threat-detection/backend/tests/test_metadata.py @@ -0,0 +1,212 @@ +""" +©AngelaMos | 2026 +test_metadata.py +""" + +import json +from pathlib import Path + +import pytest +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_sessionmaker, + create_async_engine, +) +from sqlalchemy.pool import StaticPool +from sqlalchemy import select +from sqlmodel import SQLModel + +from app.models.model_metadata import ModelMetadata +from ml.metadata import compute_model_version, save_model_metadata + + +@pytest.fixture +async def db_session(tmp_path: Path): + """ + In-memory SQLite session for metadata tests + """ + from app.models import model_metadata as _reg # noqa: F401 + + engine = create_async_engine( + "sqlite+aiosqlite://", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + async with engine.begin() as conn: + await conn.run_sync(SQLModel.metadata.create_all) + + factory = async_sessionmaker( + engine, class_=AsyncSession, expire_on_commit=False + ) + async with factory() as session: + yield session + + await engine.dispose() + + +@pytest.fixture +def model_artifacts(tmp_path: Path) -> Path: + """ + Create fake ONNX model files for version hashing + """ + (tmp_path / "ae.onnx").write_bytes(b"ae-model-data-123") + (tmp_path / "rf.onnx").write_bytes(b"rf-model-data-456") + (tmp_path / "if.onnx").write_bytes(b"if-model-data-789") + (tmp_path / "scaler.json").write_text( + json.dumps({"center": [0.0], "scale": [1.0]}) + ) + (tmp_path / "threshold.json").write_text( + json.dumps({"threshold": 0.05}) + ) + return tmp_path + + +class TestComputeModelVersion: + """ + Test SHA-256 based model version hashing + """ + + def test_returns_12_char_hex( + self, model_artifacts: Path + ) -> None: + """ + Version string is a 12-character hex digest + """ + version = compute_model_version( + model_artifacts / "ae.onnx" + ) + + assert len(version) == 12 + assert all(c in "0123456789abcdef" for c in version) + + def test_same_file_same_version( + self, model_artifacts: Path + ) -> None: + """ + Same file produces the same version string + """ + v1 = compute_model_version( + model_artifacts / "ae.onnx" + ) + v2 = compute_model_version( + model_artifacts / "ae.onnx" + ) + + assert v1 == v2 + + def test_different_files_different_versions( + self, model_artifacts: Path + ) -> None: + """ + Different files produce different version strings + """ + v_ae = compute_model_version( + model_artifacts / "ae.onnx" + ) + v_rf = compute_model_version( + model_artifacts / "rf.onnx" + ) + + assert v_ae != v_rf + + +class TestSaveModelMetadata: + """ + Test model metadata persistence to database + """ + + @pytest.mark.asyncio + async def test_creates_three_rows( + self, + db_session: AsyncSession, + model_artifacts: Path, + ) -> None: + """ + save_model_metadata creates one row per model type + """ + rows = await save_model_metadata( + db_session, + model_dir=model_artifacts, + training_samples=500, + metrics={"f1": 0.9, "pr_auc": 0.88}, + ) + + assert len(rows) == 3 + + @pytest.mark.asyncio + async def test_all_rows_active( + self, + db_session: AsyncSession, + model_artifacts: Path, + ) -> None: + """ + All newly saved models are marked as active + """ + rows = await save_model_metadata( + db_session, + model_dir=model_artifacts, + training_samples=500, + metrics={"f1": 0.9}, + ) + + assert all(r.is_active for r in rows) + + @pytest.mark.asyncio + async def test_model_types_correct( + self, + db_session: AsyncSession, + model_artifacts: Path, + ) -> None: + """ + Row model types are autoencoder, random_forest, isolation_forest + """ + rows = await save_model_metadata( + db_session, + model_dir=model_artifacts, + training_samples=500, + metrics={}, + ) + types = {r.model_type for r in rows} + + assert types == { + "autoencoder", + "random_forest", + "isolation_forest", + } + + @pytest.mark.asyncio + async def test_previous_active_replaced( + self, + db_session: AsyncSession, + model_artifacts: Path, + ) -> None: + """ + Saving new metadata replaces previous active models + """ + await save_model_metadata( + db_session, + model_dir=model_artifacts, + training_samples=500, + metrics={"f1": 0.9}, + ) + + (model_artifacts / "ae.onnx").write_bytes( + b"new-ae-data" + ) + await save_model_metadata( + db_session, + model_dir=model_artifacts, + training_samples=600, + metrics={"f1": 0.95}, + ) + + result = await db_session.execute( + select(ModelMetadata) + ) + all_rows = result.scalars().all() + active_rows = [r for r in all_rows if r.is_active] + + assert len(active_rows) == 3 + assert all( + r.training_samples == 600 for r in active_rows + )