SillyTavern_replica/backend/models/studio_models.py

"""
Studio workflow editor data models.
"""
from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field


class DisplayParam(BaseModel):
    key: str
    label: str
    type: str = "text"
    required: bool = True
    placeholder: str = ""


class InputRef(BaseModel):
    ref: str
    label: Optional[str] = None
    optional: bool = False


class ScoringDimension(BaseModel):
    id: str
    name: str
    criteria: str = ""


class InsertionRagConfig(BaseModel):
    libraryId: str = ""
    threshold: float = 0.5
    maxEntries: int = 3


class InsertionConfig(BaseModel):
    position: int = 1
    activationType: str = "permanent"
    key: str = ""
    keysecondary: str = ""
    comment: str = ""
    ragConfig: Optional[InsertionRagConfig] = None


class ScoringConfig(BaseModel):
    enabled: bool = True
    dimensions: List[ScoringDimension] = Field(default_factory=list)
    rubric: Optional[str] = None


class StudioNode(BaseModel):
    id: str
    skillId: str
    displayName: str
    enabled: bool = True
    niche: Optional[str] = None
    loopUntilSatisfied: bool = False
    config: Dict[str, Any] = Field(default_factory=dict)
    displayParams: List[DisplayParam] = Field(default_factory=list)
    inputs: List[InputRef] = Field(default_factory=list)


class PipelineDefinition(BaseModel):
    workflowGoal: str = ""
    nodes: List[StudioNode] = Field(default_factory=list)


class StudioProjectMeta(BaseModel):
    id: str
    name: str
    description: str = ""
    templateId: Optional[str] = None
    characterId: Optional[str] = None
    worldbookId: Optional[str] = None
    createdAt: str = ""
    updatedAt: str = ""


class StudioProject(BaseModel):
    meta: StudioProjectMeta
    pipeline: PipelineDefinition


class ArtifactDef(BaseModel):
    type: str
    displayName: str = ""


class SkillTemplateDef(BaseModel):
    skillId: str
    displayName: str
    description: str = ""
    displayParams: List[DisplayParam] = Field(default_factory=list)
    configWhitelist: List[str] = Field(default_factory=list)
    artifacts: List[ArtifactDef] = Field(default_factory=list)
    supportsLoopUntilSatisfied: bool = False
    supportsInputs: bool = False
    supportsInsertion: bool = False
    supportsScoring: bool = False


class WorkflowTemplateSummary(BaseModel):
    id: str
    name: str
    description: str = ""


class WorkflowVariableDef(BaseModel):
    ref: str
    label: str
    description: str = ""


class DynamicVariableSuffix(BaseModel):
    suffix: str
    labelPattern: str


class WorkflowVariablesResponse(BaseModel):
    builtIn: List[WorkflowVariableDef] = Field(default_factory=list)
    dynamic: List[WorkflowVariableDef] = Field(default_factory=list)


class SkillTemplatesCatalog(BaseModel):
    templates: List[SkillTemplateDef] = Field(default_factory=list)


class StudioProjectSummary(BaseModel):
    id: str
    name: str
    description: str = ""
    updatedAt: str = ""


class CreateStudioProjectRequest(BaseModel):
    name: str = "新项目"
    template_id: str = "builtin.studio.example"
    project_id: Optional[str] = None


class UpdateStudioProjectRequest(BaseModel):
    name: Optional[str] = Field(None, min_length=1, max_length=120)
    description: Optional[str] = Field(None, max_length=500)


class StudioRunStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class ToolQuestionOption(BaseModel):
    question: str
    options: List[str] = Field(default_factory=list)


class StepMessage(BaseModel):
    """Short step-scoped dialogue (not full chat history)."""
    id: str
    role: str  # user | assistant
    content: str
    createdAt: Optional[str] = None


class LastToolResponse(BaseModel):
    """LLM tool-call payload surfaced to the run UI (R2+)."""
    thinking: Optional[str] = None
    evaluation: Optional[str] = None
    questions: List[ToolQuestionOption] = Field(default_factory=list)
    generatedAt: Optional[str] = None


class PromptBlock(BaseModel):
    """Single assembled context section for LLM prompt (R2 debug / execution)."""
    id: str
    label: str
    content: str
    source: str = "auto"  # auto | manual | workflow


class TurnSnapshot(BaseModel):
    """State captured before each LLM turn (for undo)."""
    lastDraft: Optional[Dict[str, Any]] = None
    lastToolResponse: Optional[LastToolResponse] = None
    stepMessages: List[StepMessage] = Field(default_factory=list)
    timestamp: Optional[str] = None


class StudioNodeRunState(BaseModel):
    nodeId: str
    displayName: str
    skillId: str
    status: str  # pending | active | completed | skipped
    loopUntilSatisfied: bool = False
    lastDraft: Optional[Dict[str, Any]] = None
    lastToolResponse: Optional[LastToolResponse] = None
    stepMessages: List[StepMessage] = Field(default_factory=list)
    turnHistory: List[TurnSnapshot] = Field(default_factory=list)


class StudioRun(BaseModel):
    id: str
    projectId: str
    status: StudioRunStatus
    pipelineSnapshot: PipelineDefinition
    pipelineVersionNote: str
    currentNodeId: Optional[str] = None
    nodeStates: List[StudioNodeRunState] = Field(default_factory=list)
    workflowVariables: Dict[str, Any] = Field(default_factory=dict)
    lastPromptBlocks: List[PromptBlock] = Field(default_factory=list)
    title: str = ""
    createdAt: str = ""
    updatedAt: str = ""


class AdvanceRunRequest(BaseModel):
    displayParams: Dict[str, str] = Field(default_factory=dict)
    saveMode: Literal["advance", "append", "overwrite"] = "advance"


class SaveRunRequest(BaseModel):
    mode: Literal["incremental", "overwrite"]


class SwitchRunNodeRequest(BaseModel):
    nodeId: str = Field(..., min_length=1)


class RunMessageRequest(BaseModel):
    content: str = Field(..., min_length=1, max_length=32000)
    stream: bool = False
    profileId: Optional[str] = None
    apiConfig: Optional[Dict[str, str]] = None


class RunRerollRequest(BaseModel):
    stream: bool = False
    profileId: Optional[str] = None
    apiConfig: Optional[Dict[str, str]] = None


class RenameRunRequest(BaseModel):
    title: str = Field(..., min_length=1, max_length=120)


class StudioRunSummary(BaseModel):
    id: str
    projectId: str
    status: StudioRunStatus
    currentNodeId: Optional[str] = None
    title: str = ""
    createdAt: str = ""
    updatedAt: str = ""