SillyTavern_replica/backend/api/routes/studioRoute.py

import logging
from typing import Any, Dict, List

from fastapi import APIRouter, HTTPException

from models.studio_models import (
    AdvanceRunRequest,
    CreateStudioProjectRequest,
    PipelineDefinition,
    RenameRunRequest,
    StudioProject,
    StudioProjectSummary,
    StudioRun,
    StudioRunSummary,
    UpdateStudioProjectRequest,
    WorkflowTemplateSummary,
    WorkflowVariablesResponse,
)
from services.studio_project_service import studio_project_service
from services.studio_run_service import studio_run_service

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/studio", tags=["studio"])


@router.get("/projects", response_model=List[StudioProjectSummary])
async def list_studio_projects():
    try:
        return studio_project_service.list_projects()
    except Exception as e:
        logger.error("Failed to list studio projects: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/projects/{project_id}", response_model=StudioProject)
async def get_studio_project(project_id: str):
    try:
        return studio_project_service.get_project(project_id)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to get studio project %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.patch("/projects/{project_id}", response_model=StudioProject)
async def update_studio_project(project_id: str, req: UpdateStudioProjectRequest):
    try:
        return studio_project_service.update_project_meta(
            project_id,
            name=req.name,
            description=req.description,
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to update studio project %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.put("/projects/{project_id}/pipeline", response_model=StudioProject)
async def save_studio_pipeline(project_id: str, pipeline: PipelineDefinition):
    try:
        return studio_project_service.save_pipeline(project_id, pipeline)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to save pipeline for %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/templates", response_model=List[WorkflowTemplateSummary])
async def list_workflow_templates():
    try:
        return studio_project_service.list_workflow_templates()
    except Exception as e:
        logger.error("Failed to list workflow templates: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/variables", response_model=WorkflowVariablesResponse)
async def get_workflow_variables(projectId: str | None = None):
    try:
        return studio_project_service.get_workflow_variables(projectId)
    except Exception as e:
        logger.error("Failed to load workflow variables: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/skill-templates")
async def get_skill_templates() -> Dict[str, Any]:
    try:
        return studio_project_service.get_skill_templates()
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to load skill templates: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/niches")
async def get_niches() -> Dict[str, Any]:
    try:
        return studio_project_service.get_niches()
    except Exception as e:
        logger.error("Failed to load niches: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/projects/{project_id}")
async def delete_studio_project(project_id: str):
    try:
        studio_project_service.delete_project(project_id)
        return {"ok": True, "id": project_id}
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to delete studio project %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/projects", response_model=StudioProject)
async def create_studio_project(req: CreateStudioProjectRequest):
    try:
        return studio_project_service.create_project(req)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to create studio project: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/projects/{project_id}/runs", response_model=StudioRun)
async def create_studio_run(project_id: str):
    try:
        return studio_run_service.create_run(project_id)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to create studio run for %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/projects/{project_id}/runs", response_model=List[StudioRunSummary])
async def list_studio_runs(project_id: str):
    try:
        studio_project_service.get_project(project_id)
        return studio_run_service.list_runs(project_id)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to list studio runs for %s: %s", project_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/projects/{project_id}/runs/{run_id}", response_model=StudioRun)
async def get_studio_run(project_id: str, run_id: str):
    try:
        return studio_run_service.get_run(project_id, run_id)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error("Failed to get studio run %s/%s: %s", project_id, run_id, e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/projects/{project_id}/runs/{run_id}/advance", response_model=StudioRun)
async def advance_studio_run(
    project_id: str, run_id: str, req: AdvanceRunRequest
):
    try:
        return studio_run_service.advance_run(
            project_id, run_id, display_params=req.displayParams
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except NotImplementedError as e:
        raise HTTPException(status_code=501, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(
            "Failed to advance studio run %s/%s: %s", project_id, run_id, e
        )
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/projects/{project_id}/runs/{run_id}")
async def delete_studio_run(project_id: str, run_id: str):
    try:
        studio_run_service.delete_run(project_id, run_id)
        return {"ok": True, "id": run_id}
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(
            "Failed to delete studio run %s/%s: %s", project_id, run_id, e
        )
        raise HTTPException(status_code=500, detail=str(e))


@router.patch("/projects/{project_id}/runs/{run_id}", response_model=StudioRun)
async def rename_studio_run(
    project_id: str, run_id: str, req: RenameRunRequest
):
    try:
        return studio_run_service.rename_run(project_id, run_id, req.title)
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(
            "Failed to rename studio run %s/%s: %s", project_id, run_id, e
        )
        raise HTTPException(status_code=500, detail=str(e))