SillyTavern_replica/backend/services/workflow_engine.py

"""
Workflow engine – orchestrates template loading, state machine execution, and run persistence.
"""
from __future__ import annotations

import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Awaitable, Callable, Dict, List, Optional

try:
    from backend.core.config import settings
    from backend.models.agent import (
        ChatRunBinding,
        ChatTurnResult,
        RunEvent,
        RunStatus,
        TurnCallbacks,
        TurnContext,
        WorkflowRun,
        WorkflowTemplate,
        WorkflowTemplateKind,
    )
    from backend.services.state_machine_runner import StateMachineRunner
    from backend.services.tool_registry import ToolRegistry, default_tool_registry
    from backend.services.tools.chat_tools import register_chat_tools
except ImportError:
    from core.config import settings
    from models.agent import (
        ChatRunBinding,
        ChatTurnResult,
        RunEvent,
        RunStatus,
        TurnCallbacks,
        TurnContext,
        WorkflowRun,
        WorkflowTemplate,
        WorkflowTemplateKind,
    )
    from services.state_machine_runner import StateMachineRunner
    from services.tool_registry import ToolRegistry, default_tool_registry
    from services.tools.chat_tools import register_chat_tools


class WorkflowEngine:
    def __init__(self, registry: Optional[ToolRegistry] = None) -> None:
        self.registry = registry or default_tool_registry
        if not self.registry.list_specs():
            register_chat_tools(self.registry)

    def _template_dir(self, template_id: str) -> Path:
        return settings.AGENT_TEMPLATES_PATH / template_id

    def load_template(self, template_id: str = WorkflowTemplateKind.BUILTIN_CHAT.value) -> WorkflowTemplate:
        template_path = self._template_dir(template_id) / "template.json"
        with open(template_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return WorkflowTemplate(**data)

    def _run_dir(self, role_name: str, chat_name: str) -> Path:
        return settings.AGENT_RUNS_PATH / "chat" / role_name / chat_name

    def _persist_run(self, run: WorkflowRun, events: List[RunEvent]) -> None:
        run_dir = self._run_dir(run.binding.role_name, run.binding.chat_name)
        run_dir.mkdir(parents=True, exist_ok=True)

        run_file = run_dir / "run.json"
        run.finished_at = datetime.now().isoformat()
        with open(run_file, "w", encoding="utf-8") as f:
            json.dump(run.model_dump(), f, ensure_ascii=False, indent=2)

        events_file = run_dir / "events.jsonl"
        with open(events_file, "a", encoding="utf-8") as f:
            for event in events:
                f.write(json.dumps(event.model_dump(), ensure_ascii=False) + "\n")

    async def run_turn(
        self,
        request_data: Dict[str, Any],
        *,
        stream: bool = False,
        on_chunk: Optional[Callable[[str], Awaitable[None]]] = None,
        on_worldbook_active: Optional[Callable[[List[Any]], Awaitable[None]]] = None,
        on_tasks_created: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None,
        template_id: str = WorkflowTemplateKind.BUILTIN_CHAT.value,
    ) -> ChatTurnResult:
        current_role = request_data.get("currentRole", "")
        current_chat = request_data.get("currentChat", "")
        user_message = request_data.get("mes", "")

        if not current_role or not user_message:
            return ChatTurnResult(
                success=False,
                error="缺少必要的参数：currentRole 或 mes",
                workflow_template_id=template_id,
            )

        preset_config = request_data.get("presetConfig", {})
        preset_name = preset_config.get("selectedPreset")

        run_id = uuid.uuid4().hex
        binding = ChatRunBinding(
            role_name=current_role,
            chat_name=current_chat or "",
            template_id=template_id,
        )
        run = WorkflowRun(
            id=run_id,
            template_id=template_id,
            binding=binding,
            status=RunStatus.PENDING,
        )

        callbacks = TurnCallbacks(
            on_chunk=on_chunk,
            on_worldbook_active=on_worldbook_active,
            on_tasks_created=on_tasks_created,
        )

        ctx = TurnContext(
            request_data=request_data,
            template_id=template_id,
            run_id=run_id,
            stream=stream,
            callbacks=callbacks,
            current_role=current_role,
            current_chat=current_chat or "",
            user_message=user_message,
            preset_name=preset_name,
        )

        template = self.load_template(template_id)
        sm_path = self._template_dir(template_id) / template.state_machine_path
        runner = StateMachineRunner.from_file(sm_path, self.registry)

        try:
            events = await runner.run(run, ctx)
            self._persist_run(run, events)

            active_entries = [
                entry.model_dump() if hasattr(entry, "model_dump") else entry
                for entry in ctx.active_entries
            ]

            return ChatTurnResult(
                success=True,
                content=ctx.generated_content,
                active_entries=active_entries,
                task_ids=ctx.task_ids,
                run_id=run_id,
                workflow_template_id=template_id,
            )
        except Exception as exc:
            run.status = RunStatus.FAILED
            run.error = str(exc)
            try:
                self._persist_run(run, [])
            except Exception:
                pass
            return ChatTurnResult(
                success=False,
                error=f"工作流执行失败: {exc}",
                run_id=run_id,
                workflow_template_id=template_id,
            )


# Module-level singleton
workflow_engine = WorkflowEngine()