SillyTavern_replica/backend/services/studio_run_service.py

"""
Create and load Studio pipeline runs with frozen pipeline snapshots.
"""
from __future__ import annotations

import copy
import json
import logging
import shutil
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from core.config import settings
from models.studio_models import (
    PipelineDefinition,
    StudioNode,
    StudioNodeRunState,
    StudioRun,
    StudioRunStatus,
    StudioRunSummary,
)
from services.character_service import CharacterService
from services.studio_context_service import store_context_on_run
from services.studio_project_service import studio_project_service
from services.worldbook_service import worldbook_service

logger = logging.getLogger(__name__)


def _read_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _write_json(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _build_node_states(pipeline: PipelineDefinition) -> tuple[list[StudioNodeRunState], Optional[str]]:
    """First enabled node is active; other enabled nodes pending; disabled skipped."""
    states: list[StudioNodeRunState] = []
    first_active_id: Optional[str] = None
    seen_active = False

    for node in pipeline.nodes:
        if not node.enabled:
            status = "skipped"
        elif not seen_active:
            status = "active"
            first_active_id = node.id
            seen_active = True
        else:
            status = "pending"

        states.append(
            StudioNodeRunState(
                nodeId=node.id,
                displayName=node.displayName,
                skillId=node.skillId,
                status=status,
                loopUntilSatisfied=node.loopUntilSatisfied,
            )
        )

    return states, first_active_id


def _find_node(pipeline: PipelineDefinition, node_id: str) -> Optional[StudioNode]:
    for node in pipeline.nodes:
        if node.id == node_id:
            return node
    return None


def _next_enabled_node_id(pipeline: PipelineDefinition, after_node_id: str) -> Optional[str]:
    seen = False
    for node in pipeline.nodes:
        if node.id == after_node_id:
            seen = True
            continue
        if seen and node.enabled:
            return node.id
    return None


class StudioRunService:
    def __init__(self) -> None:
        self._character_service = CharacterService()

    @property
    def runs_root(self) -> Path:
        return settings.AGENT_STUDIO_RUNS_PATH

    def _project_runs_dir(self, project_id: str) -> Path:
        return self.runs_root / project_id

    def _run_dir(self, project_id: str, run_id: str) -> Path:
        return self._project_runs_dir(project_id) / run_id

    def _run_path(self, project_id: str, run_id: str) -> Path:
        return self._run_dir(project_id, run_id) / "run.json"

    def _save_run(self, project_id: str, run_id: str, run: StudioRun) -> None:
        _write_json(
            self._run_path(project_id, run_id),
            run.model_dump(mode="json"),
        )

    def create_run(self, project_id: str) -> StudioRun:
        project = studio_project_service.get_project(project_id)
        snapshot = PipelineDefinition(**copy.deepcopy(project.pipeline.model_dump()))
        now = datetime.now().isoformat()
        run_id = str(uuid.uuid4())
        node_states, current_node_id = _build_node_states(snapshot)

        has_active = current_node_id is not None
        run = StudioRun(
            id=run_id,
            projectId=project_id,
            status=StudioRunStatus.RUNNING if has_active else StudioRunStatus.PENDING,
            pipelineSnapshot=snapshot,
            pipelineVersionNote=now,
            currentNodeId=current_node_id,
            nodeStates=node_states,
            workflowVariables={"workflow.goal": snapshot.workflowGoal},
            createdAt=now,
            updatedAt=now,
        )
        self._save_run(project_id, run_id, run)
        return run

    def list_runs(self, project_id: str) -> List[StudioRunSummary]:
        root = self._project_runs_dir(project_id)
        if not root.exists():
            return []

        summaries: List[StudioRunSummary] = []
        for child in sorted(root.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True):
            if not child.is_dir():
                continue
            run_path = child / "run.json"
            if not run_path.exists():
                continue
            raw = _read_json(run_path)
            summaries.append(
                StudioRunSummary(
                    id=raw.get("id", child.name),
                    projectId=raw.get("projectId", project_id),
                    status=StudioRunStatus(raw.get("status", StudioRunStatus.PENDING.value)),
                    currentNodeId=raw.get("currentNodeId"),
                    title=raw.get("title", ""),
                    createdAt=raw.get("createdAt", ""),
                    updatedAt=raw.get("updatedAt", ""),
                )
            )
        return summaries

    def get_run(self, project_id: str, run_id: str) -> StudioRun:
        run_path = self._run_path(project_id, run_id)
        if not run_path.exists():
            raise FileNotFoundError(f"Studio run not found: {project_id}/{run_id}")
        run = StudioRun(**_read_json(run_path))
        return run

    def _validate_display_params(
        self, node: StudioNode, display_params: Dict[str, str]
    ) -> Dict[str, str]:
        normalized: Dict[str, str] = {}
        for dp in node.displayParams:
            raw = display_params.get(dp.key, "")
            value = (raw or "").strip()
            if dp.required and not value:
                raise ValueError(f"缺少必填项：{dp.label}")
            normalized[dp.key] = value
        return normalized

    def _execute_init_bind(
        self,
        project_id: str,
        run: StudioRun,
        node: StudioNode,
        display_params: Dict[str, str],
    ) -> tuple[Dict[str, Any], Dict[str, Any]]:
        params = self._validate_display_params(node, display_params)
        character_name = params.get("characterName", "")
        worldbook_name = params.get("worldbookName", "")

        if self._character_service.get_character_by_name(character_name):
            raise ValueError(f"角色「{character_name}」已存在")
        if worldbook_service._get_worldbook_path(worldbook_name).exists():
            raise ValueError(f"世界书「{worldbook_name}」已存在")

        worldbook = worldbook_service.create_worldbook(worldbook_name)
        worldbook_id = worldbook["id"]

        character = self._character_service.create_character(
            {
                "name": character_name,
                "description": "",
                "personality": "",
                "scenario": "",
                "first_mes": "",
                "mes_example": "",
                "categories": [],
                "tags": [],
                "worldInfoId": worldbook_id,
            }
        )
        character_id = character.id

        studio_project_service.update_project_bindings(
            project_id, character_id, worldbook_id
        )

        workflow_variables = dict(run.workflowVariables or {})
        workflow_variables["workflow.goal"] = run.pipelineSnapshot.workflowGoal
        workflow_variables["workflow.boundCharacter"] = (
            f"名称：{character_name}\nID：{character_id}"
        )
        workflow_variables["workflow.boundWorldbook"] = (
            f"名称：{worldbook_name}\nID：{worldbook_id}"
        )

        last_draft = {
            "displayParams": params,
            "characterId": character_id,
            "worldbookId": worldbook_id,
            "characterName": character_name,
            "worldbookName": worldbook_name,
        }
        return workflow_variables, last_draft

    def advance_run(
        self,
        project_id: str,
        run_id: str,
        display_params: Optional[Dict[str, str]] = None,
    ) -> StudioRun:
        run = self.get_run(project_id, run_id)
        if run.status not in (StudioRunStatus.RUNNING, StudioRunStatus.PENDING):
            raise ValueError("运行已结束，无法继续推进")

        current_node_id = run.currentNodeId
        if not current_node_id:
            raise ValueError("当前运行无活动节点")

        current_node = _find_node(run.pipelineSnapshot, current_node_id)
        if not current_node:
            raise ValueError(f"节点不存在：{current_node_id}")

        current_state = next(
            (s for s in run.nodeStates if s.nodeId == current_node_id), None
        )
        if not current_state or current_state.status != "active":
            raise ValueError("当前节点不可执行")

        workflow_variables = dict(run.workflowVariables or {})
        last_draft: Optional[Dict[str, Any]] = None

        if current_node.skillId == "studio.init_bind":
            if not display_params:
                raise ValueError("请填写引导表单后再提交")
            workflow_variables, last_draft = self._execute_init_bind(
                project_id, run, current_node, display_params
            )
        else:
            raise NotImplementedError(
                f"技能「{current_node.skillId}」的执行尚未实现（R2+）"
            )

        next_node_id = _next_enabled_node_id(run.pipelineSnapshot, current_node_id)
        now = datetime.now().isoformat()

        new_node_states: list[StudioNodeRunState] = []
        for state in run.nodeStates:
            updated = state.model_copy()
            if state.nodeId == current_node_id:
                updated.status = "completed"
                if last_draft is not None:
                    updated.lastDraft = last_draft
            elif next_node_id and state.nodeId == next_node_id:
                updated.status = "active"
            new_node_states.append(updated)

        new_status = (
            StudioRunStatus.COMPLETED if not next_node_id else StudioRunStatus.RUNNING
        )
        updated_run = run.model_copy(
            update={
                "currentNodeId": next_node_id,
                "nodeStates": new_node_states,
                "workflowVariables": workflow_variables,
                "status": new_status,
                "updatedAt": now,
            }
        )
        updated_run = store_context_on_run(updated_run, next_node_id)
        self._save_run(project_id, run_id, updated_run)
        return updated_run

    def delete_run(self, project_id: str, run_id: str) -> None:
        run_dir = self._run_dir(project_id, run_id)
        if not run_dir.exists():
            raise FileNotFoundError(f"Studio run not found: {project_id}/{run_id}")
        shutil.rmtree(run_dir)

    def rename_run(self, project_id: str, run_id: str, title: str) -> StudioRun:
        run = self.get_run(project_id, run_id)
        trimmed = (title or "").strip()
        if not trimmed:
            raise ValueError("运行名称不能为空")
        now = datetime.now().isoformat()
        updated = run.model_copy(update={"title": trimmed, "updatedAt": now})
        self._save_run(project_id, run_id, updated)
        return updated


studio_run_service = StudioRunService()