Update storage/media_routers.py

storage/media_routers.py

@@ -17,6 +17,8 @@ router = APIRouter(prefix="/media", tags=["Media Manager"])
 MEDIA_ROOT = Path("/data/media")
 file_manager = FileManager(MEDIA_ROOT)
 HF_TOKEN = os.getenv("HF_TOKEN")
+VALID_SUBTYPES = ("HITL", "MoE", "Salamandra")
+AUDIO_EXTENSIONS = ["*.mp3", "*.wav", "*.aac", "*.m4a", "*.ogg", "*.flac"]
 
 
 @router.delete("/clear_media", tags=["Media Manager"])
@@ -537,3 +539,132 @@ def list_all_audios(
         })
 
     return results
+
+
+@router.post("/upload_audio_subtype", tags=["Media Manager"])
+async def upload_audio_subtype(
+    audio: UploadFile = File(...),
+    sha1: str = Query(..., description="SHA1 of the video folder"),
+    subtype: str = Query(..., description="Subtype: HITL, MoE, or Salamandra"),
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    Upload audio for a given subtype (HITL, MoE, Salamandra).
+    - Creates folder if missing: /data/media/<sha1>/<subtype>/
+    - Deletes any previous audio files
+    - Saves the new audio
+    """
+    validate_token(token)
+
+    if subtype not in VALID_SUBTYPES:
+        raise HTTPException(status_code=400, detail=f"subtype must be one of {VALID_SUBTYPES}")
+
+    MEDIA_ROOT = Path("/data/media")
+    subtype_dir = MEDIA_ROOT / sha1 / subtype
+    subtype_dir.mkdir(parents=True, exist_ok=True)
+
+    # Delete old audio files
+    try:
+        for pattern in AUDIO_EXTENSIONS:
+            for old_audio in subtype_dir.glob(pattern):
+                old_audio.unlink()
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to delete old audio files: {exc}")
+
+    final_path = subtype_dir / audio.filename
+
+    try:
+        file_bytes = await audio.read()
+        with open(final_path, "wb") as f:
+            f.write(file_bytes)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to save audio: {exc}")
+
+    return JSONResponse(
+        status_code=200,
+        content={
+            "status": "ok",
+            "sha1": sha1,
+            "subtype": subtype,
+            "saved_to": str(final_path)
+        }
+    )
+
+
+@router.get("/download_audio_subtype", tags=["Media Manager"])
+def download_audio_subtype(
+    sha1: str,
+    subtype: str,
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    Download the first audio file for a given subtype.
+    """
+    validate_token(token)
+
+    if subtype not in VALID_SUBTYPES:
+        raise HTTPException(status_code=400, detail=f"subtype must be one of {VALID_SUBTYPES}")
+
+    MEDIA_ROOT = Path("/data/media")
+    subtype_dir = MEDIA_ROOT / sha1 / subtype
+
+    if not subtype_dir.exists() or not subtype_dir.is_dir():
+        raise HTTPException(status_code=404, detail=f"{subtype} folder not found")
+
+    # Find audio files
+    audio_files = []
+    for pattern in AUDIO_EXTENSIONS:
+        audio_files.extend(list(subtype_dir.glob(pattern)))
+
+    if not audio_files:
+        raise HTTPException(status_code=404, detail="No audio files found")
+
+    audio_path = audio_files[0]
+
+    return FileResponse(
+        path=audio_path,
+        media_type="audio/" + audio_path.suffix.lstrip("."),
+        filename=audio_path.name
+    )
+
+
+@router.get("/list_subtype_audios", tags=["Media Manager"])
+def list_subtype_audios(
+    sha1: str = Query(..., description="SHA1 of the video folder"),
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    List the most recent audio file for each subtype (HITL, MoE, Salamandra)
+    under /data/media/<sha1>.
+
+    Returns:
+    - sha1: folder name
+    - subtype: name of the subtype
+    - audio_name: latest audio file or None
+    """
+    validate_token(token)
+
+    results = []
+
+    for subtype in VALID_SUBTYPES:
+        subtype_dir = MEDIA_ROOT / sha1 / subtype
+
+        latest_audio = None
+
+        if subtype_dir.exists() and subtype_dir.is_dir():
+            files = []
+            for pattern in AUDIO_EXTENSIONS:
+                files.extend(list(subtype_dir.glob(pattern)))
+
+            if files:
+                # Sort by modification time (newest first)
+                files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
+                latest_audio = files[0].name
+
+        results.append({
+            "sha1": sha1,
+            "subtype": subtype,
+            "audio_name": latest_audio
+        })
+
+    return results