SillyTavern_replica/backend/services/studio_context_service.py

"""
Assemble Studio run prompt context from pipeline snapshot, workflow variables,
and node outputs (R2). Does not include full chat history.
"""
from __future__ import annotations

import json
import re
from typing import Any, Dict, List, Optional

from models.studio_models import (
    PipelineDefinition,
    PromptBlock,
    StudioNode,
    StudioNodeRunState,
    StudioRun,
)

_NODE_OUTPUT_REF = re.compile(r"^([^.]+)\.output$")

AUTO_BLOCK_SPECS = (
    ("currentProduct", "目前产物", "auto"),
    ("thinkingFlow", "思考流程", "auto"),
    ("coreGoal", "核心目的", "auto"),
    ("scoringCriteria", "评价标准与优化建议", "auto"),
)


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 _state_map(run: StudioRun) -> Dict[str, StudioNodeRunState]:
    return {s.nodeId: s for s in run.nodeStates}


def _format_draft(draft: Optional[Dict[str, Any]]) -> str:
    if not draft:
        return "（暂无内容）"
    for key in ("entryContent", "content", "text", "body"):
        value = draft.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return json.dumps(draft, ensure_ascii=False, indent=2)


def _format_scoring(config: Dict[str, Any]) -> str:
    scoring = config.get("scoring") or {}
    if not scoring.get("enabled", True):
        return "（本步骤未启用评价）"
    dimensions = scoring.get("dimensions") or []
    if not dimensions:
        rubric = scoring.get("rubric")
        if rubric:
            return str(rubric).strip()
        return "（未配置评价维度）"
    lines: List[str] = []
    for dim in dimensions:
        name = dim.get("name") or dim.get("id") or "维度"
        criteria = (dim.get("criteria") or "").strip()
        lines.append(f"- {name}：{criteria}" if criteria else f"- {name}")
    return "\n".join(lines)


def _resolve_workflow_ref(ref: str, workflow_variables: Dict[str, Any]) -> str:
    value = workflow_variables.get(ref)
    if value is None:
        return "（尚未可用）"
    if isinstance(value, str):
        return value.strip() or "（空）"
    return json.dumps(value, ensure_ascii=False, indent=2)


def _resolve_node_output_ref(
    ref: str,
    pipeline: PipelineDefinition,
    state_by_id: Dict[str, StudioNodeRunState],
) -> str:
    match = _NODE_OUTPUT_REF.match(ref)
    if not match:
        return f"（无法解析引用：{ref}）"
    node_id = match.group(1)
    source_node = _find_node(pipeline, node_id)
    source_state = state_by_id.get(node_id)
    label = source_node.displayName if source_node else node_id
    if not source_state or source_state.status != "completed":
        return f"（前序步骤「{label}」尚未完成）"
    return _format_draft(source_state.lastDraft)


def _auto_block_content(
    block_id: str,
    node: StudioNode,
    node_state: Optional[StudioNodeRunState],
) -> str:
    config = node.config or {}
    if block_id == "currentProduct":
        return _format_draft(node_state.lastDraft if node_state else None)
    if block_id == "thinkingFlow":
        return (config.get("thinkingPrompt") or "").strip() or "（未配置思考流程）"
    if block_id == "coreGoal":
        return (config.get("stepGoal") or "").strip() or "（未配置步骤目标）"
    if block_id == "scoringCriteria":
        return _format_scoring(config)
    return ""


def assemble_prompt_blocks(run: StudioRun, node_id: str) -> List[PromptBlock]:
    """
    Build ordered prompt blocks for a worldbook step from inputs[].ref,
    workflow variables, node outputs, and auto-injected context items.
    """
    pipeline = run.pipelineSnapshot
    node = _find_node(pipeline, node_id)
    if not node:
        return []

    if node.skillId != "studio.worldbook_entry":
        return []

    state_by_id = _state_map(run)
    node_state = state_by_id.get(node_id)
    workflow_variables = dict(run.workflowVariables or {})
    blocks: List[PromptBlock] = []
    seen_ids: set[str] = set()

    def append_block(
        block_id: str,
        label: str,
        content: str,
        source: str,
    ) -> None:
        if block_id in seen_ids:
            return
        seen_ids.add(block_id)
        blocks.append(
            PromptBlock(
                id=block_id,
                label=label,
                content=content,
                source=source,
            )
        )

    for inp in node.inputs or []:
        ref = (inp.ref or "").strip()
        if not ref:
            continue
        label = (inp.label or ref).strip()
        block_id = f"ref:{ref}"

        if ref.startswith("workflow."):
            content = _resolve_workflow_ref(ref, workflow_variables)
            if inp.optional and content in ("（尚未可用）", "（空）"):
                continue
            append_block(block_id, label, content, "workflow")
            continue

        if _NODE_OUTPUT_REF.match(ref):
            content = _resolve_node_output_ref(ref, pipeline, state_by_id)
            if inp.optional and content.startswith("（前序步骤"):
                continue
            append_block(block_id, label, content, "manual")
            continue

        append_block(block_id, label, f"（未知引用类型：{ref}）", "manual")

    for block_id, label, source in AUTO_BLOCK_SPECS:
        content = _auto_block_content(block_id, node, node_state)
        append_block(block_id, label, content, source)

    return blocks


def store_context_on_run(run: StudioRun, node_id: Optional[str]) -> StudioRun:
    """Attach assembled prompt blocks to run for debug / frontend display."""
    if not node_id:
        return run.model_copy(update={"lastPromptBlocks": []})

    node = _find_node(run.pipelineSnapshot, node_id)
    if not node or node.skillId != "studio.worldbook_entry":
        return run.model_copy(update={"lastPromptBlocks": []})

    blocks = assemble_prompt_blocks(run, node_id)
    return run.model_copy(update={"lastPromptBlocks": blocks})