Merge pull request #1835 from transformerlab/fix/notes-revamp

Revamp experiment notes fully

Author: deep1401

api/test/test_notes_helpers.py

@@ -0,0 +1,150 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+
+def _make_open_mock(content: str):
+    """Return a mock async context manager that reads `content`."""
+    mock_file = AsyncMock()
+    mock_file.read = AsyncMock(return_value=content)
+    mock_file.write = AsyncMock()
+    cm = AsyncMock()
+    cm.__aenter__ = AsyncMock(return_value=mock_file)
+    cm.__aexit__ = AsyncMock(return_value=False)
+    return cm, mock_file
+
+
+def test_read_notes_returns_new_path_when_exists():
+    """read_notes returns notes/readme.md content when it exists."""
+
+    async def _run():
+        from transformerlab.routers.experiment.notes import read_notes
+
+        open_cm, mock_file = _make_open_mock("# my notes")
+        with (
+            patch("transformerlab.routers.experiment.notes.Experiment") as mock_exp_cls,
+            patch("transformerlab.routers.experiment.notes.storage") as mock_storage,
+        ):
+            mock_exp = MagicMock()
+            mock_exp.get_dir = AsyncMock(return_value="/experiments/exp1")
+            mock_exp_cls.return_value = mock_exp
+            mock_storage.join.side_effect = lambda *parts: "/".join(parts)
+            mock_storage.open = AsyncMock(return_value=open_cm)
+            result = await read_notes("exp1")
+        assert result == "# my notes"
+        mock_storage.open.assert_called_once_with("/experiments/exp1/notes/readme.md", "r", encoding="utf-8")
+
+    asyncio.run(_run())
+
+
+def test_read_notes_falls_back_to_legacy():
+    """read_notes falls back to root readme.md when notes/readme.md is missing."""
+
+    async def _run():
+        from transformerlab.routers.experiment.notes import read_notes
+
+        legacy_cm, mock_file = _make_open_mock("# legacy notes")
+        call_count = 0
+
+        async def open_side_effect(path, mode, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise FileNotFoundError
+            return legacy_cm
+
+        with (
+            patch("transformerlab.routers.experiment.notes.Experiment") as mock_exp_cls,
+            patch("transformerlab.routers.experiment.notes.storage") as mock_storage,
+        ):
+            mock_exp = MagicMock()
+            mock_exp.get_dir = AsyncMock(return_value="/experiments/exp1")
+            mock_exp_cls.return_value = mock_exp
+            mock_storage.join.side_effect = lambda *parts: "/".join(parts)
+            mock_storage.open.side_effect = open_side_effect
+            result = await read_notes("exp1")
+        assert result == "# legacy notes"
+
+    asyncio.run(_run())
+
+
+def test_read_notes_returns_empty_when_neither_exists():
+    """read_notes returns empty string when both paths are missing."""
+
+    async def _run():
+        from transformerlab.routers.experiment.notes import read_notes
+
+        async def open_always_raises(path, mode, **kwargs):
+            raise FileNotFoundError
+
+        with (
+            patch("transformerlab.routers.experiment.notes.Experiment") as mock_exp_cls,
+            patch("transformerlab.routers.experiment.notes.storage") as mock_storage,
+        ):
+            mock_exp = MagicMock()
+            mock_exp.get_dir = AsyncMock(return_value="/experiments/exp1")
+            mock_exp_cls.return_value = mock_exp
+            mock_storage.join.side_effect = lambda *parts: "/".join(parts)
+            mock_storage.open.side_effect = open_always_raises
+            result = await read_notes("exp1")
+        assert result == ""
+
+    asyncio.run(_run())
+
+
+def test_migrate_if_needed_moves_legacy_file():
+    """migrate_if_needed copies readme.md → notes/readme.md and removes the original."""
+
+    async def _run():
+        from transformerlab.routers.experiment.notes import migrate_if_needed
+
+        read_cm, _ = _make_open_mock("# old notes")
+        write_cm, write_file = _make_open_mock("")
+
+        open_calls = []
+
+        async def open_side_effect(path, mode, **kwargs):
+            open_calls.append((path, mode))
+            if mode == "r":
+                return read_cm
+            return write_cm
+
+        with (
+            patch("transformerlab.routers.experiment.notes.Experiment") as mock_exp_cls,
+            patch("transformerlab.routers.experiment.notes.storage") as mock_storage,
+        ):
+            mock_exp = MagicMock()
+            mock_exp.get_dir = AsyncMock(return_value="/experiments/exp1")
+            mock_exp_cls.return_value = mock_exp
+            mock_storage.join.side_effect = lambda *parts: "/".join(parts)
+            mock_storage.exists = AsyncMock(side_effect=[False, True])  # notes missing, legacy exists
+            mock_storage.makedirs = AsyncMock()
+            mock_storage.open.side_effect = open_side_effect
+            mock_storage.rm = AsyncMock()
+            await migrate_if_needed("exp1")
+        mock_storage.makedirs.assert_called_once()
+        write_file.write.assert_called_once_with("# old notes")
+        mock_storage.rm.assert_called_once_with("/experiments/exp1/readme.md")
+
+    asyncio.run(_run())
+
+
+def test_migrate_if_needed_skips_when_notes_already_exist():
+    """migrate_if_needed does nothing when notes/readme.md already exists."""
+
+    async def _run():
+        from transformerlab.routers.experiment.notes import migrate_if_needed
+
+        with (
+            patch("transformerlab.routers.experiment.notes.Experiment") as mock_exp_cls,
+            patch("transformerlab.routers.experiment.notes.storage") as mock_storage,
+        ):
+            mock_exp = MagicMock()
+            mock_exp.get_dir = AsyncMock(return_value="/experiments/exp1")
+            mock_exp_cls.return_value = mock_exp
+            mock_storage.join.side_effect = lambda *parts: "/".join(parts)
+            mock_storage.exists = AsyncMock(return_value=True)  # notes already exists
+            mock_storage.rm = AsyncMock()
+            await migrate_if_needed("exp1")
+        mock_storage.rm.assert_not_called()
+
+    asyncio.run(_run())

api/transformerlab/routers/experiment/experiment.py

@@ -11,6 +11,7 @@
 from transformerlab.routers.experiment import (
     documents,
     jobs,
+    notes,
     task as task_router,
 )
 from transformerlab.routers.auth import get_user_and_team
@@ -39,6 +40,12 @@
     tags=["task"],
     dependencies=[Depends(require_permission("experiment", "read", id_param="experimentId"))],
 )
+router.include_router(
+    router=notes.router,
+    prefix="/{experimentId}",
+    tags=["notes"],
+    dependencies=[Depends(require_permission("experiment", "read", id_param="experimentId"))],
+)
 
 
 @router.get("/", summary="Get all Experiments", tags=["experiment"])

api/transformerlab/routers/experiment/notes.py

@@ -0,0 +1,152 @@
+import logging
+import os
+from typing import Annotated
+
+from fastapi import APIRouter, Body, Depends, HTTPException, UploadFile
+from fastapi.responses import Response
+from werkzeug.utils import secure_filename
+
+from lab import Experiment, storage
+from transformerlab.services.permission_service import require_permission
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/notes", tags=["notes"])
+
+ALLOWED_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".svg"}
+MAX_ASSET_BYTES = 20 * 1024 * 1024
+
+IMAGE_MEDIA_TYPES = {
+    ".png": "image/png",
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".gif": "image/gif",
+    ".svg": "image/svg+xml",
+}
+
+
+async def read_notes(experimentId: str) -> str:
+    """Read notes content, falling back to legacy readme.md at experiment root."""
+    exp_obj = Experiment(experimentId)
+    experiment_dir = await exp_obj.get_dir()
+    notes_file = storage.join(experiment_dir, "notes", "readme.md")
+
+    try:
+        async with await storage.open(notes_file, "r", encoding="utf-8") as f:
+            return await f.read()
+    except FileNotFoundError:
+        pass
+
+    legacy_file = storage.join(experiment_dir, "readme.md")
+    try:
+        async with await storage.open(legacy_file, "r", encoding="utf-8") as f:
+            return await f.read()
+    except FileNotFoundError:
+        return ""
+
+
+async def migrate_if_needed(experimentId: str) -> None:
+    """Move legacy readme.md → notes/readme.md on first save if needed."""
+    exp_obj = Experiment(experimentId)
+    experiment_dir = await exp_obj.get_dir()
+    notes_dir = storage.join(experiment_dir, "notes")
+    notes_file = storage.join(notes_dir, "readme.md")
+    legacy_file = storage.join(experiment_dir, "readme.md")
+
+    if await storage.exists(notes_file):
+        return
+
+    if not await storage.exists(legacy_file):
+        return
+
+    async with await storage.open(legacy_file, "r", encoding="utf-8") as f:
+        content = await f.read()
+
+    await storage.makedirs(notes_dir, exist_ok=True)
+
+    async with await storage.open(notes_file, "w", encoding="utf-8") as f:
+        await f.write(content)
+
+    await storage.rm(legacy_file)
+
+
+@router.get("")
+async def get_notes(experimentId: str):
+    return await read_notes(experimentId)
+
+
+@router.post("")
+async def save_notes(
+    experimentId: str,
+    file_contents: Annotated[str, Body()],
+    _: None = Depends(require_permission("experiment", "write", id_param="experimentId")),
+):
+    await migrate_if_needed(experimentId)
+    exp_obj = Experiment(experimentId)
+    experiment_dir = await exp_obj.get_dir()
+    notes_dir = storage.join(experiment_dir, "notes")
+    notes_file = storage.join(notes_dir, "readme.md")
+    await storage.makedirs(notes_dir, exist_ok=True)
+    # Write a space instead of empty string — some storage backends treat empty files as deletions
+    async with await storage.open(notes_file, "w", encoding="utf-8") as f:
+        await f.write(file_contents if file_contents.strip() else " ")
+    return {"message": "OK"}
+
+
+@router.post("/assets")
+async def upload_asset(
+    experimentId: str,
+    file: UploadFile,
+    _: None = Depends(require_permission("experiment", "write", id_param="experimentId")),
+):
+    filename = secure_filename(file.filename or "")
+    if not filename:
+        raise HTTPException(status_code=400, detail="Invalid filename")
+
+    _, ext = os.path.splitext(filename.lower())
+    if ext not in ALLOWED_IMAGE_EXTENSIONS:
+        raise HTTPException(
+            status_code=400,
+            detail=f'File type "{ext}" not allowed. Allowed: {", ".join(sorted(ALLOWED_IMAGE_EXTENSIONS))}',
+        )
+
+    content = await file.read()
+    if len(content) > MAX_ASSET_BYTES:
+        raise HTTPException(status_code=400, detail="File exceeds 20 MB limit")
+
+    exp_obj = Experiment(experimentId)
+    experiment_dir = await exp_obj.get_dir()
+    assets_dir = storage.join(experiment_dir, "notes", "assets")
+    await storage.makedirs(assets_dir, exist_ok=True)
+
+    asset_path = storage.join(assets_dir, filename)
+    async with await storage.open(asset_path, "wb") as f:
+        await f.write(content)
+
+    return {"path": f"notes/assets/{filename}"}
+
+
+@router.get("/assets/{filename}")
+async def get_asset(experimentId: str, filename: str):
+    filename = secure_filename(filename)
+    if not filename:
+        raise HTTPException(status_code=400, detail="Invalid filename")
+
+    _, ext = os.path.splitext(filename.lower())
+    if ext not in ALLOWED_IMAGE_EXTENSIONS:
+        raise HTTPException(status_code=400, detail="File type not allowed")
+
+    exp_obj = Experiment(experimentId)
+    experiment_dir = await exp_obj.get_dir()
+    asset_path = storage.join(experiment_dir, "notes", "assets", filename)
+
+    if not await storage.exists(asset_path):
+        raise HTTPException(status_code=404, detail=f"Asset '{filename}' not found")
+
+    async with await storage.open(asset_path, "rb") as f:
+        content = await f.read()
+
+    return Response(
+        content=content,
+        media_type=IMAGE_MEDIA_TYPES.get(ext, "application/octet-stream"),
+    )

cli/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab-cli"
-version = "0.0.23"
+version = "0.0.24"
 description = "Transformer Lab CLI"
 requires-python = ">=3.10"
 authors = [{ name = "Transformer Lab", email = "hello@transformerlab.ai" }]